From 32691794f9eab704b233814896de056b6234ead0 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Fri, 23 Oct 2026 10:58:36 +0200 Subject: [PATCH 1/2] rebase --- AGENTS.md | 165 +++++++++++---- core/src/main/resources/python/jax_helper.py | 18 ++ core/src/main/scala/dimwit/package.scala | 1 + .../tensor/tensorops/ControlFlowOps.scala | 189 ++++++++++++++++++ .../tensor/tensorops/FunctionalOps.scala | 23 ++- .../tensor/TensorOpsControlFlowSuite.scala | 144 +++++++++++++ .../scala/dimwit/basic/SIRCalibration.scala | 81 ++++++++ .../scala/dimwit/basic/SIRSimulation.scala | 129 ++++++------ mdocs/AGENTS.md | 62 ++++++ 9 files changed, 697 insertions(+), 115 deletions(-) create mode 100644 core/src/main/scala/dimwit/tensor/tensorops/ControlFlowOps.scala create mode 100644 core/src/test/scala/dimwit/tensor/TensorOpsControlFlowSuite.scala create mode 100644 examples/src/main/scala/dimwit/basic/SIRCalibration.scala diff --git a/AGENTS.md b/AGENTS.md index ec79c480..b5ef95f2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -893,6 +893,83 @@ println(s"Reduced A: ${summedA}") println(s"Reduced B: ${summedB}") ``` +### Loops and Branches: scan, foriLoop, whileLoop, cond + +A Scala loop inside `jit` is unrolled while tracing: 100 iterations become 100 copies of the body, +and the compile time grows with them. `scan`, `foriLoop`, `whileLoop` and `cond` wrap the `jax.lax` +primitives of the same names, whose body is traced once. They work eagerly and under `jit`. + +The state carried from one iteration to the next may be any tensor tree (a tensor, a tuple, a case class). +Since the body is traced once, it must not change the shape of the carry. + +```scala +import dimwit.* + +trait Time derives Label +trait Feature derives Label + +val signal = Tensor1(Axis[Time]).fromArray(Array(1.0f, 3.0f, 2.0f, 6.0f)) + +// scan: loop over the axis Time of a tensor. The body gets the carry and the slice at each +// position along Time, and returns the next carry and an output; the outputs are stacked along Time. +val (lastAverage, averages) = scan(Axis[Time])(Tensor0(0.0f), signal): (average, x) => + val next = average * 0.5f + x * 0.5f + (next, next) +// averages: Tensor1[Time, Float32] = [0.5, 1.75, 1.875, 3.9375] + +// scan over several tensors along the same axis, like zipvmap; () means no outputs +val weights = Tensor1(Axis[Time]).fromArray(Array(1.0f, 1.0f, 2.0f, 2.0f)) +val (weightedSum, _) = scan(Axis[Time])(Tensor0(0.0f), (signal, weights)): + case (sum, (x, weight)) => (sum + x * weight, ()) + +// foriLoop: loop over the indices lower until upper. The index is a traced Tensor0[Int32], +// e.g. to write a window at it +val squares = foriLoop(0, 4)(Tensor(Shape(Axis[Time] -> 4)).fill(0)): (i, buffer) => + buffer.set(Axis[Time].at(i, 1))((i * i).broadcastTo(Shape(Axis[Time] -> 1))) +// squares: [0, 1, 4, 9] + +// whileLoop: loop while a condition holds, here Newton's method for the square root of 2 +val (root, iterations) = whileLoop((Tensor0(1.0f), Tensor0(0)))((x, _) => (x * x - Tensor0(2.0f)).abs > Tensor0(1e-5f)): + case (x, n) => ((x + Tensor0(2.0f) / x) * 0.5f, n + Tensor0(1)) + +// cond: choose a branch by a predicate that may be traced, unlike Scala's if; +// both branches must return the same shapes +val relu = jit((x: Tensor0[Float32]) => cond(x > Tensor0(0.0f))(x)(Tensor0(0.0f))) +``` + +`scan` and `foriLoop` are differentiable, so `Autodiff.grad` through them gives the same gradient as through the +unrolled loop. JAX cannot reverse-differentiate `whileLoop`, since the number of iterations is only known at runtime; +use forward mode (`Autodiff.jacFwd`) for it instead. + +```scala +// ERROR: the scanned axis must exist in the tensor +val notAnAxis = scan(Axis[Feature])(Tensor0(0.0f), signal)((sum, x) => (sum + x, ())) +// error: +// Axis[MdocApp5.this.Feature] not found in Tensor[Tuple1[MdocApp5.this.Time]]. +// I found: +// +// dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[Tuple1[MdocApp5.this.Time], +// MdocApp5.this.Feature, R]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp5.this.Time, +// EmptyTuple.type, MdocApp5.this.Feature]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A, B, L]), +// ???) +// +// But given instance concatRight in object AxisIndex does not match type dimwit.tensor.ShapeTypeHelpers.AxisIndex[EmptyTuple.type, MdocApp5.this.Feature]. +// val notAnAxis = scan(Axis[Feature])(Tensor0(0.0f), signal)((sum, x) => (sum + x, ())) +// ^ +``` + +```scala +// ERROR at runtime: the body must not change the shape of the carry. +// (The error is raised inside the traced body; here only its message is shown.) +val shapeError = scala.util.Try: + scan(Axis[Time])(Tensor1(Axis[Feature]).fromArray(Array(0.0f)), signal): (state, x) => + (concatenate(state, state, Axis[Feature]), ()) +.failed.get.getMessage.linesIterator.map(_.trim).find(_.contains("requirement failed")).get +// shapeError: String = "java.lang.IllegalArgumentException: requirement failed: The body of a loop must not change the shape of the carry, but the result is Shape(Feature -> 1) of Float32 and then Shape(Feature -> 2) of Float32" +``` + --- ## Automatic Differentiation @@ -1321,11 +1398,11 @@ trait D derives Label val intTensor = Tensor1(Axis[A]).fromArray(Array(1, 2, 3)) val wrong = intTensor.exp // exp requires IsFloating constraint // error: -// value exp is not a member of dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.tensor.DType.Int32]. +// value exp is not a member of dimwit.tensor.Tensor1[MdocApp13.this.A, dimwit.tensor.DType.Int32]. // An extension method was tried, but could not be fully constructed: // -// dimwit.exp[Tuple1[MdocApp12.this.A], dimwit.tensor.DType.Int32](this.intTensor)( -// dimwit.tensor.Labels.consTuple[MdocApp12.this.A, EmptyTuple.type]( +// dimwit.exp[Tuple1[MdocApp13.this.A], dimwit.tensor.DType.Int32](this.intTensor)( +// dimwit.tensor.Labels.consTuple[MdocApp13.this.A, EmptyTuple.type]( // this.A.derived$Label, dimwit.tensor.Labels.emptyTuple), // /* missing */ // summon[dimwit.tensor.ValueTypeClasses.IsFloating[dimwit.tensor.DType.Int32]] @@ -1341,12 +1418,12 @@ val wrong = intTensor.exp // exp requires IsFloating constraint val boolTensor = Tensor1(Axis[A]).fromArray(Array(true, false, true)) val wrong = boolTensor.mean // error: -// value mean is not a member of dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.tensor.DType.Bool]. +// value mean is not a member of dimwit.tensor.Tensor1[MdocApp13.this.A, dimwit.tensor.DType.Bool]. // An extension method was tried, but could not be fully constructed: // -// dimwit.mean[Tuple1[MdocApp12.this.A], dimwit.tensor.DType.Bool](this.boolTensor) +// dimwit.mean[Tuple1[MdocApp13.this.A], dimwit.tensor.DType.Bool](this.boolTensor) // ( -// dimwit.tensor.Labels.consTuple[MdocApp12.this.A, EmptyTuple.type]( +// dimwit.tensor.Labels.consTuple[MdocApp13.this.A, EmptyTuple.type]( // this.A.derived$Label, dimwit.tensor.Labels.emptyTuple), // /* missing */ // summon[dimwit.tensor.ValueTypeClasses.IsFloating[dimwit.tensor.DType.Bool]] @@ -1367,9 +1444,9 @@ val t1 = Tensor1(Axis[A]).fromArray(Array(1.0f, 2.0f)) val t2 = Tensor1(Axis[B]).fromArray(Array(3.0f, 4.0f, 5.0f)) val wrong = t1 + t2 // Different labels AND different sizes // error: -// Found: (MdocApp12.this.t2 : -// dimwit.tensor.Tensor1[MdocApp12.this.B, dimwit.tensor.DType.Float32]) -// Required: dimwit.tensor.Tensor[Tuple1[MdocApp12.this.A], dimwit.tensor.DType.Float32] +// Found: (MdocApp13.this.t2 : +// dimwit.tensor.Tensor1[MdocApp13.this.B, dimwit.tensor.DType.Float32]) +// Required: dimwit.tensor.Tensor[Tuple1[MdocApp13.this.A], dimwit.tensor.DType.Float32] ``` ```scala @@ -1378,22 +1455,22 @@ val m1 = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) // Shape val m2 = Tensor2(Axis[C], Axis[D]).fromArray(Array(Array(3.0f), Array(4.0f))) // Shape: (2, 1) val wrong = m1.dot(Axis[B])(m2) // Axis[B] not in m2 // error: -// Axis[MdocApp12.this.B] not found in Tensor[(MdocApp12.this.C, MdocApp12.this.D)]. +// Axis[MdocApp13.this.B] not found in Tensor[(MdocApp13.this.C, MdocApp13.this.D)]. // I found: // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ -// (MdocApp12.this.C, MdocApp12.this.D), MdocApp12.this.B, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp12.this.C, -// MdocApp12.this.D *: EmptyTuple.type, MdocApp12.this.B]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp12.this.D, -// EmptyTuple.type, MdocApp12.this.B]( +// (MdocApp13.this.C, MdocApp13.this.D), MdocApp13.this.B, R]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp13.this.C, +// MdocApp13.this.D *: EmptyTuple.type, MdocApp13.this.B]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp13.this.D, +// EmptyTuple.type, MdocApp13.this.B]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A, B², L]) // ), // ???) // -// But given instance concatRight in object AxisIndex does not match type dimwit.tensor.ShapeTypeHelpers.AxisIndex[EmptyTuple.type, MdocApp12.this.B] +// But given instance concatRight in object AxisIndex does not match type dimwit.tensor.ShapeTypeHelpers.AxisIndex[EmptyTuple.type, MdocApp13.this.B] // -// where: B is a trait in class MdocApp12 +// where: B is a trait in class MdocApp13 // B² is a type variable with constraint <: Tuple // . ``` @@ -1406,7 +1483,7 @@ val t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) val wrong = t + 10.0f // Should use +! for scalar broadcast // error: // Found: (10.0f : Float) -// Required: dimwit.tensor.Tensor[(MdocApp12.this.A, MdocApp12.this.B), +// Required: dimwit.tensor.Tensor[(MdocApp13.this.A, MdocApp13.this.B), // dimwit.tensor.DType.Float32] ``` @@ -1417,31 +1494,31 @@ val t2 = Tensor1(Axis[A]).fromArray(Array(3.0f, 4.0f)) // This works but is semantically wrong (use + instead) val wrong = t1 +! t2 // error: -// Cannot broadcast tensors of shapes Tuple1[MdocApp12.this.A] and Tuple1[MdocApp12.this.A]. If same shape no broadcasting allowed!. +// Cannot broadcast tensors of shapes Tuple1[MdocApp13.this.A] and Tuple1[MdocApp13.this.A]. If same shape no broadcasting allowed!. // I found: // -// dimwit.tensor.Broadcast.broadcastLeft[Tuple1[MdocApp12.this.A], -// Tuple1[MdocApp12.this.A], dimwit.tensor.DType.Float32]( -// dimwit.tensor.Labels.consTuple[MdocApp12.this.A, EmptyTuple.type]( +// dimwit.tensor.Broadcast.broadcastLeft[Tuple1[MdocApp13.this.A], +// Tuple1[MdocApp13.this.A], dimwit.tensor.DType.Float32]( +// dimwit.tensor.Labels.consTuple[MdocApp13.this.A, EmptyTuple.type]( // this.A.derived$Label, dimwit.tensor.Labels.emptyTuple), -// dimwit.tensor.Labels.consTuple[MdocApp12.this.A, EmptyTuple.type]( +// dimwit.tensor.Labels.consTuple[MdocApp13.this.A, EmptyTuple.type]( // this.A.derived$Label, dimwit.tensor.Labels.emptyTuple), -// dimwit.tensor.TupleHelpers.StrictSubset.bridge[Tuple1[MdocApp12.this.A], -// Tuple1[MdocApp12.this.A]]( -// dimwit.tensor.TupleHelpers.Subset.consTuple²[MdocApp12.this.A, -// EmptyTuple.type, Tuple1[MdocApp12.this.A]]( -// dimwit.tensor.TupleHelpers.SetMember.found[MdocApp12.this.A, +// dimwit.tensor.TupleHelpers.StrictSubset.bridge[Tuple1[MdocApp13.this.A], +// Tuple1[MdocApp13.this.A]]( +// dimwit.tensor.TupleHelpers.Subset.consTuple²[MdocApp13.this.A, +// EmptyTuple.type, Tuple1[MdocApp13.this.A]]( +// dimwit.tensor.TupleHelpers.SetMember.found[MdocApp13.this.A, // EmptyTuple.type], -// dimwit.tensor.TupleHelpers.Subset.emptyTuple²[Tuple1[MdocApp12.this.A]]), +// dimwit.tensor.TupleHelpers.Subset.emptyTuple²[Tuple1[MdocApp13.this.A]]), // /* missing */ // summon[ -// scala.util.NotGiven[Tuple1[MdocApp12.this.A] =:= -// Tuple1[MdocApp12.this.A]] +// scala.util.NotGiven[Tuple1[MdocApp13.this.A] =:= +// Tuple1[MdocApp13.this.A]] // ] // ) // ) // -// But no implicit values were found that match type scala.util.NotGiven[Tuple1[MdocApp12.this.A] =:= Tuple1[MdocApp12.this.A]] +// But no implicit values were found that match type scala.util.NotGiven[Tuple1[MdocApp13.this.A] =:= Tuple1[MdocApp13.this.A]] // // where: consTuple is a given instance in object Labels // consTuple² is a given instance in object Subset @@ -1459,24 +1536,24 @@ val wrong = t1 +! t2 val t = Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.0f, 2.0f))) val wrong = t.sum(Axis[C]) // Axis[C] not in tensor // error: -// Axis[MdocApp12.this.C] not found in Tensor[(MdocApp12.this.A, MdocApp12.this.B)]. +// Axis[MdocApp13.this.C] not found in Tensor[(MdocApp13.this.A, MdocApp13.this.B)]. // I found: // // dimwit.tensor.ShapeTypeHelpers.AxisRemover.bridge[ -// (MdocApp12.this.A, MdocApp12.this.B), MdocApp12.this.C, R]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp12.this.A, -// MdocApp12.this.B *: EmptyTuple.type, MdocApp12.this.C]( -// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp12.this.B, -// EmptyTuple.type, MdocApp12.this.C]( +// (MdocApp13.this.A, MdocApp13.this.B), MdocApp13.this.C, R]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp13.this.A, +// MdocApp13.this.B *: EmptyTuple.type, MdocApp13.this.C]( +// dimwit.tensor.ShapeTypeHelpers.AxisIndex.search[MdocApp13.this.B, +// EmptyTuple.type, MdocApp13.this.C]( // dimwit.tensor.ShapeTypeHelpers.AxisIndex.concatRight[A², B², L]) // ), // ???) // -// But given instance concatRight in object AxisIndex does not match type dimwit.tensor.ShapeTypeHelpers.AxisIndex[EmptyTuple.type, MdocApp12.this.C] +// But given instance concatRight in object AxisIndex does not match type dimwit.tensor.ShapeTypeHelpers.AxisIndex[EmptyTuple.type, MdocApp13.this.C] // -// where: A is a trait in class MdocApp12 +// where: A is a trait in class MdocApp13 // A² is a type variable with constraint <: Tuple -// B is a trait in class MdocApp12 +// B is a trait in class MdocApp13 // B² is a type variable with constraint <: Tuple // . ``` @@ -1486,7 +1563,7 @@ val wrong = t.sum(Axis[C]) // Axis[C] not in tensor val t = Tensor2(Axis[A], Axis[B]).fill(1.0f) val wrong = t.vmap(Axis[C])(_.sum) // Axis[C] doesn't exist // error: -// value fill is not a member of dimwit.tensor.Tensor2.Axes2Factory[MdocApp12.this.A, MdocApp12.this.B] +// value fill is not a member of dimwit.tensor.Tensor2.Axes2Factory[MdocApp13.this.A, MdocApp13.this.B] ``` ### Gradient Errors @@ -1528,8 +1605,8 @@ val wrong = Autodiff.grad(nonScalar) // Use jacobian instead // t2Tree²: dimwit.tensortree.TensorTree[T2²], outTree³: // dimwit.tensortree.TensorTree[dimwit.tensor.Tensor0[V³]]): (T1², T2²) => // dimwit.autodiff.Grad[(T1², T2²)] -// match arguments (dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.Float32] => -// dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.Float32]) +// match arguments (dimwit.tensor.Tensor1[MdocApp13.this.A, dimwit.Float32] => +// dimwit.tensor.Tensor1[MdocApp13.this.A, dimwit.Float32]) // // where: T1 is a type variable // T1² is a type variable diff --git a/core/src/main/resources/python/jax_helper.py b/core/src/main/resources/python/jax_helper.py index fe1b7618..6a04c31d 100644 --- a/core/src/main/resources/python/jax_helper.py +++ b/core/src/main/resources/python/jax_helper.py @@ -85,3 +85,21 @@ def jit(f): def jit_fn(f, jit_kwargs=None): return wrap(jax.jit, f, kwargs=jit_kwargs) + + +# Structured control flow. The body is traced once, so a loop does not unroll under jit. +# Each Scala callback is wrapped in a Python lambda, which shields it from JAX introspection. + +def scan(f, init, xs): + """jax.lax.scan; f(carry, x) returns the Python tuple (carry, y).""" + return jax.lax.scan(lambda carry, x: tuple(f(carry, x)), init, xs) + +def fori_loop(lower, upper, body, init): + return jax.lax.fori_loop(lower, upper, lambda i, carry: body(i, carry), init) + +def while_loop(cond, body, init): + return jax.lax.while_loop(lambda carry: cond(carry), lambda carry: body(carry), init) + +def cond(pred, true_fn, false_fn): + """jax.lax.cond without operands: the branches are closures.""" + return jax.lax.cond(pred, lambda: true_fn(None), lambda: false_fn(None)) diff --git a/core/src/main/scala/dimwit/package.scala b/core/src/main/scala/dimwit/package.scala index a5d33c0e..75fb9271 100644 --- a/core/src/main/scala/dimwit/package.scala +++ b/core/src/main/scala/dimwit/package.scala @@ -89,6 +89,7 @@ package object dimwit: export dimwit.tensor.tensorops.ElementWiseOps.{maximum, minimum, maximum_!, minimum_!} export dimwit.tensor.tensorops.StructuralOps.{where, where_!, triu, tril, stack, concatenate} export dimwit.tensor.tensorops.FunctionalOps.zipvmap + export dimwit.tensor.tensorops.ControlFlowOps.{scan, foriLoop, whileLoop, cond} // Export convolution options export dimwit.tensor.{Padding, Stride1, Stride2, Stride3} diff --git a/core/src/main/scala/dimwit/tensor/tensorops/ControlFlowOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/ControlFlowOps.scala new file mode 100644 index 00000000..9bb43fe0 --- /dev/null +++ b/core/src/main/scala/dimwit/tensor/tensorops/ControlFlowOps.scala @@ -0,0 +1,189 @@ +package dimwit.tensor.tensorops + +import dimwit.OnError +import dimwit.jax.Jax +import dimwit.python.PyIndex.itemAt +import dimwit.tensor.Axis +import dimwit.tensor.DType.Bool +import dimwit.tensor.DType.Int32 +import dimwit.tensor.Label +import dimwit.tensor.Labels +import dimwit.tensor.ShapeTypeHelpers.AxisRemover +import dimwit.tensor.Tensor +import dimwit.tensor.Tensor0 +import dimwit.tensor.ShapeTypeHelpers.SharedAxisRemover +import dimwit.tensor.tensorops.FunctionalOps.PrependAxis +import dimwit.tensor.tensorops.FunctionalOps.ZipVmap.ShapesOf +import dimwit.tensor.tensorops.FunctionalOps.ZipVmap.TensorsOf +import dimwit.tensor.tensorops.FunctionalOps.ZipVmap.ValuesOf +import dimwit.tensor.tensorops.FunctionalOps.ZipVmap.slicesOf +import dimwit.tensortree.TensorTree +import me.shadaj.scalapy.py +import me.shadaj.scalapy.py.SeqConverters + +/** Structured control flow, e.g. `scan` or `whileLoop`. The body of a loop is traced once, + * so that a loop does not unroll under `jit`, however many iterations it runs. + */ +private[dimwit] object ControlFlowOps: + + /** Loops over the axis `L` of `xs`, like `jax.lax.scan`: the body `f` gets the carry and the slice of `xs` + * at each position along `L`, and returns the next carry and an output. The outputs are stacked along `L`. + * + * The body is traced once, so it may not change the shape of the carry. It works under `jit` and is differentiable. + * + * @param axis The axis to loop over. + * @param init The initial carry, any tensor tree. + * @param xs The tensor to loop over. + * @param f The body, from the carry and the slice of `xs` to the next carry and an output. + * @return The final carry, and the outputs with `L` prepended to every tensor. + */ + def scan[L: Label, Carry, T <: Tuple, V, Y](axis: Axis[L])(init: Carry, xs: Tensor[T, V])(using + ev: AxisRemover[T, L], + sliceLabels: Labels[ev.RemainingAxes] + )(f: (Carry, Tensor[ev.RemainingAxes, V]) => (Carry, Y))(using + carryTree: TensorTree[Carry], + yTree: TensorTree[Y], + prependAxis: PrependAxis[L, Y], + ysTree: TensorTree[prependAxis.Out] + ): (Carry, prependAxis.Out) = + scanAlong[L, Carry, Tensor[ev.RemainingAxes, V], Y, prependAxis.Out]( + init, + Jax.jnp.moveaxis(xs.jaxValue, ev.index, 0), + List(xs.shape.dimensions(ev.index)) + )(Tensor[ev.RemainingAxes, V](_))(f) + + /** Loops over the axis `L` of all tensors of `xs` together, like `jax.lax.scan`: the body `f` gets the carry and + * the tuple of their slices at each position along `L`, and returns the next carry and an output. + * The outputs are stacked along `L`. + * + * The body is traced once, so it may not change the shape of the carry. It works under `jit` and is differentiable. + * + * @param axis The axis to loop over. Every tensor of `xs` must have it, with the same extent. + * @param init The initial carry, any tensor tree. + * @param xs The tuple of tensors to loop over. + * @param f The body, from the carry and the tuple of slices of `xs` to the next carry and an output. + * @return The final carry, and the outputs with `L` prepended to every tensor. + */ + def scan[L: Label, Carry, Inputs <: Tuple, Y](axis: Axis[L])(init: Carry, xs: Inputs)(using + ev: SharedAxisRemover[ShapesOf[Inputs], L] + )(f: (Carry, TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]]) => (Carry, Y))(using + carryTree: TensorTree[Carry], + yTree: TensorTree[Y], + prependAxis: PrependAxis[L, Y], + ysTree: TensorTree[prependAxis.Out] + ): (Carry, prependAxis.Out) = + val tensors = xs.toList.map(_.asInstanceOf[Tensor[?, ?]]) + scanAlong[L, Carry, TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]], Y, prependAxis.Out]( + init, + py.Dynamic.global.tuple(tensors.zip(ev.indices).map((t, i) => Jax.jnp.moveaxis(t.jaxValue, i, 0)).toPythonProxy), + tensors.zip(ev.indices).map((t, i) => t.shape.dimensions(i)) + )(slices => slicesOf[ShapesOf[Inputs], L, ValuesOf[Inputs]](slices.as[Seq[Jax.PyDynamic]], ev))(f) + + /** The scan over the axis `L` of the tensors with the extents `extents`, moved to the front in the pytree `leading`. */ + private def scanAlong[L: Label, Carry, X, Y, Ys](init: Carry, leading: py.Any, extents: List[Int])( + sliceOf: Jax.PyDynamic => X + )(f: (Carry, X) => (Carry, Y))(using + carryTree: TensorTree[Carry], + yTree: TensorTree[Y], + ysTree: TensorTree[Ys] + ): (Carry, Ys) = + val axisName = summon[Label[L]].name + require( + extents.distinct.sizeIs == 1, + s"All tensors scanned over Axis[$axisName] must have the same extent along it, but they have ${extents.mkString(", ")}" + ) + val body = (carryPy: Jax.PyDynamic, xPy: Jax.PyDynamic) => + OnError.traceStack: + val carry = carryTree.fromPyTree(carryPy) + val (next, y) = f(carry, sliceOf(xPy)) + requireSameShapes(carry, next) + // the Python tuple (carry, y) that jax.lax.scan expects + py.Dynamic.global.tuple(Seq(carryTree.toPyTree(next), yTree.toPyTree(y)).toPythonProxy): py.Any + val result = Jax.jax_helper.scan(body, carryTree.toPyTree(init), leading) + (carryTree.fromPyTree(result.itemAt(0)), ysTree.fromPyTree(result.itemAt(1))) + + /** Loops over the indices `lower until upper`, like `jax.lax.fori_loop`: the body gets the index and the carry, + * and returns the next carry. + * + * The index is a traced tensor, e.g. to slice or set at it. The body is traced once, so it may not change + * the shape of the carry. Since the bounds are static, the loop works under `jit` and is differentiable. + * + * @param lower The first index. + * @param upper The index after the last one. + * @param init The initial carry, any tensor tree. + * @param body The body, from the index and the carry to the next carry. + * @return The final carry. + */ + def foriLoop[Carry](lower: Int, upper: Int)(init: Carry)(body: (Tensor0[Int32], Carry) => Carry)(using + carryTree: TensorTree[Carry] + ): Carry = + val bodyPy = (indexPy: Jax.PyDynamic, carryPy: Jax.PyDynamic) => + OnError.traceStack: + val carry = carryTree.fromPyTree(carryPy) + val next = body(Tensor[EmptyTuple, Int32](indexPy), carry) + requireSameShapes(carry, next) + carryTree.toPyTree(next) + carryTree.fromPyTree(Jax.jax_helper.fori_loop(lower, upper, bodyPy, carryTree.toPyTree(init))) + + /** Loops while `condition` holds, like `jax.lax.while_loop`: the body gets the carry and returns the next carry. + * + * The number of iterations may depend on the values, e.g. to stop once a computation has converged. + * The body is traced once, so it may not change the shape of the carry. + * JAX cannot reverse-differentiate a while loop, so `Autodiff.grad` through it fails; forward mode (`Autodiff.jacFwd`) works. + * + * @param init The initial carry, any tensor tree. + * @param condition Whether to run the body (again) on the carry. + * @param body The body, from the carry to the next carry. + * @return The first carry for which `condition` does not hold. + */ + def whileLoop[Carry](init: Carry)(condition: Carry => Tensor0[Bool])(body: Carry => Carry)(using + carryTree: TensorTree[Carry] + ): Carry = + val conditionPy = (carryPy: Jax.PyDynamic) => + OnError.traceStack: + condition(carryTree.fromPyTree(carryPy)).jaxValue + val bodyPy = (carryPy: Jax.PyDynamic) => + OnError.traceStack: + val carry = carryTree.fromPyTree(carryPy) + val next = body(carry) + requireSameShapes(carry, next) + carryTree.toPyTree(next) + carryTree.fromPyTree(Jax.jax_helper.while_loop(conditionPy, bodyPy, carryTree.toPyTree(init))) + + /** Evaluates `ifTrue` or `ifFalse` depending on `predicate`, like `jax.lax.cond`. + * Unlike Scala's `if`, the predicate may be traced, e.g. under `jit` or in a loop body. + * Both branches are traced, so they must return the same shapes. + * + * @param predicate Which branch to evaluate. + * @param ifTrue The result if `predicate` holds, any tensor tree. + * @param ifFalse The result otherwise. + */ + def cond[T](predicate: Tensor0[Bool])(ifTrue: => T)(ifFalse: => T)(using tree: TensorTree[T]): T = + // JAX traces both branches, one after the other; the second one is checked against the first one + var firstBranch: Option[T] = None + def branch(value: => T) = (_: Jax.PyDynamic) => + OnError.traceStack: + val result = value + firstBranch match + case None => firstBranch = Some(result) + case Some(first) => requireSameShapes(first, result, "The branches of cond must return the same shapes") + tree.toPyTree(result) + tree.fromPyTree(Jax.jax_helper.cond(predicate.jaxValue, branch(ifTrue), branch(ifFalse))) + + /** Requires that the tensor trees `before` and `after` have the same shapes and value types, naming the first tensor that differs. */ + private def requireSameShapes[T](before: T, after: T, problem: String = "The body of a loop must not change the shape of the carry")(using + tree: TensorTree[T] + ): Unit = + leaves(before).zip(leaves(after)).foreach: + case ((path, b), (_, a)) => + val name = if path.isEmpty then "the result" else s"`$path`" + require( + b.shape.dimensions == a.shape.dimensions && b.dtype == a.dtype, + s"$problem, but $name is ${b.shape} of ${b.dtype} and then ${a.shape} of ${a.dtype}" + ) + + /** The tensors of the tree `t`, with their paths (e.g. `cache.keys`). */ + private def leaves[T](t: T)(using tree: TensorTree[T]): List[(String, Tensor[?, ?])] = + val result = List.newBuilder[(String, Tensor[?, ?])] + tree.foreachWithName(t, [S <: Tuple, V] => (_: Labels[S]) ?=> (path: String, tensor: Tensor[S, V]) => result += ((path, tensor)): Unit) + result.result() diff --git a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala index 83ac9212..5fc9a382 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala @@ -25,12 +25,12 @@ import scala.annotation.implicitNotFound private[dimwit] object FunctionalOps: /** Prepends the axis `L` to every tensor of the tensor tree `FOut`: the result - * type of a `vmap`/`zipvmap` whose body returned `FOut`. + * type of a `vmap`/`zipvmap` whose body returned `FOut`, or the outputs of a `scan`. * - * Instances cover tensors, tuples, named tuples and any nesting of those. + * Instances cover tensors, tuples, named tuples, any nesting of those, and `Unit` (no output). */ @implicitNotFound( - "Cannot prepend Axis[${L}] to ${FOut}. A vmap/zipvmap body may return a Tensor, a tuple, a named tuple, or any nesting of those" + "Cannot prepend Axis[${L}] to ${FOut}. A vmap/zipvmap/scan body may return a Tensor, a tuple, a named tuple, or any nesting of those" ) trait PrependAxis[L, FOut]: type Out @@ -45,6 +45,8 @@ private[dimwit] object FunctionalOps: given tensor[L, Shape <: Tuple, V]: Aux[L, Tensor[Shape, V], Tensor[L *: Shape, V]] = instance + given unit[L]: Aux[L, Unit, Unit] = instance + given emptyTuple[L]: Aux[L, EmptyTuple, EmptyTuple] = instance given consTuple[L, H, HOut, T <: Tuple, TOut <: Tuple](using @@ -71,6 +73,15 @@ private[dimwit] object FunctionalOps: type ShapesOf[Tensors <: Tuple] = Tuple.Map[Tensors, ExtractShape] type ValuesOf[Tensors <: Tuple] = Tuple.Map[Tensors, ExtractValue] + /** The JAX arrays `slices` as the tuple of tensors they are: slices along `L` of tensors with the shapes `Shapes`. */ + private[dimwit] def slicesOf[Shapes <: Tuple, L, Values <: Tuple]( + slices: Seq[Jax.PyDynamic], + ev: SharedAxisRemover[Shapes, L] + ): TensorsOf[ev.RemainingAxes, Values] = + val tensors = slices.zip(ev.shapesLabels).map: (jaxArr, labels) => + Tensor(jaxArr)(using LabelsImpl(labels)) + Tuple.fromArray(tensors.toArray).asInstanceOf[TensorsOf[ev.RemainingAxes, Values]] + /** Zips the given given tensors along the specified axis * and applies the function `f` to the zipped tensors. * @@ -103,11 +114,7 @@ private[dimwit] object FunctionalOps: ): prependAxis.Out = val fpy = (args: py.Dynamic) => OnError.traceStack: - val tensorList = args.as[Seq[py.Dynamic]].zip(ev.shapesLabels).map: (jaxArr, labels) => - Tensor(jaxArr)(using LabelsImpl(labels)) - - val inputTuple = Tuple.fromArray(tensorList.toArray) - val result = f(inputTuple.asInstanceOf[TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]]]) + val result = f(slicesOf[ShapesOf[Inputs], L, ValuesOf[Inputs]](args.as[Seq[py.Dynamic]], ev)) toPyTree.toPyTree(result) val jaxInputs = py.Dynamic.global.tuple(tensors.toArray.map(_.asInstanceOf[Tensor[?, ?]].jaxValue).toPythonProxy) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsControlFlowSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsControlFlowSuite.scala new file mode 100644 index 00000000..a21857f9 --- /dev/null +++ b/core/src/test/scala/dimwit/tensor/TensorOpsControlFlowSuite.scala @@ -0,0 +1,144 @@ +package dimwit.tensor + +import dimwit.* +import dimwit.jax.Jax +import me.shadaj.scalapy.py + +object TensorOpsControlFlowSuite: + trait Time derives Label + trait Feature derives Label + + case class Cache(keys: Tensor2[Time, Feature, Float32], count: Tensor0[Int32]) + +class TensorOpsControlFlowSuite extends DimwitTest: + + import TensorOpsControlFlowSuite.* + + // Shape: Time=4, Feature=3, the value at (t, f) is t + f / 10 + val xs = Tensor2(Axis[Time], Axis[Feature]).fromArray(Array.tabulate(4, 3)((t, f) => t + f / 10f)) + val zeros = Tensor1(Axis[Feature]).fromArray(Array(0f, 0f, 0f)) + + /** The running sum of `xs` along Time, computed with a Scala fold. */ + def runningSum(xs: Tensor2[Time, Feature, Float32]): (Tensor1[Feature, Float32], Tensor2[Time, Feature, Float32]) = + val sums = xs.unstack(Axis[Time]).scanLeft(zeros)(_ + _).tail + (sums.last, stack(sums, Axis[Time])) + + /** The number of operations in the jaxpr of `f`, without those inside loop bodies. */ + def jaxprSize(f: Tensor1[Feature, Float32] => Tensor1[Feature, Float32]): Int = + val fpy = (x: py.Dynamic) => f(Tensor[Tuple1[Feature], Float32](x)).jaxValue + py.Dynamic.global.len(Jax.jax.make_jaxpr(Jax.jax_helper.wrap_fn(fpy))(zeros.jaxValue).jaxpr.eqns).as[Int] + + describe("scan over an axis"): + + it("gives the same result as a fold"): + val (sum, sums) = scan(Axis[Time])(zeros, xs)((carry, x) => (carry + x, carry + x)) + val (expectedSum, expectedSums) = runningSum(xs) + sum should approxEqual(expectedSum) + sums should approxEqual(expectedSums) + + it("gives the same result under jit"): + val jitScan = jit((xs: Tensor2[Time, Feature, Float32]) => scan(Axis[Time])(zeros, xs)((carry, x) => (carry + x, carry + x))) + val (sum, sums) = jitScan(xs) + val (expectedSum, expectedSums) = runningSum(xs) + sum should approxEqual(expectedSum) + sums should approxEqual(expectedSums) + + it("scans over an axis that is not the first one and prepends it to the outputs"): + val (sum, sums) = scan(Axis[Time])(zeros, xs.transpose)((carry, x) => (carry + x, carry + x)) + val (expectedSum, expectedSums) = runningSum(xs) + sum should approxEqual(expectedSum) + sums.axes shouldBe List("Time", "Feature") + sums should approxEqual(expectedSums) + + it("scans over a tuple of tensors with a case class carry"): + val weights = Tensor1(Axis[Time]).fromArray(Array(1f, 2f, 3f, 4f)) + val init = Cache(Tensor(Shape(Axis[Time] -> 1, Axis[Feature] -> 3)).fill(0f), Tensor0(0)) + val (cache, counts) = scan(Axis[Time])(init, (xs, weights)): + case (cache, (x, weight)) => + val keys = cache.keys + (x *! weight).prependAxis(Axis[Time]) + (Cache(keys, cache.count + Tensor0(1)), cache.count) + val weightedSum = xs.unstack(Axis[Time]).zip(Seq(1f, 2f, 3f, 4f)).map((x, w) => x *! w).reduce(_ + _) + cache.count.item shouldBe 4 + cache.keys.slice(Axis[Time].at(0)) should approxEqual(weightedSum) + counts shouldEqual Tensor1(Axis[Time]).fromArray(Array(0, 1, 2, 3)) + + it("has the same gradient as the unrolled loop"): + def scanned(w: Tensor1[Feature, Float32]): Tensor0[Float32] = + scan(Axis[Time])(zeros, xs)((h, x) => ((h * w + x).tanh, ()))._1.sum + def unrolled(w: Tensor1[Feature, Float32]): Tensor0[Float32] = + xs.unstack(Axis[Time]).foldLeft(zeros)((h, x) => (h * w + x).tanh).sum + val w = Tensor1(Axis[Feature]).fromArray(Array(0.5f, -0.3f, 0.8f)) + scanned(w) should approxEqual(unrolled(w)) + Autodiff.grad(scanned)(w).value should approxEqual(Autodiff.grad(unrolled)(w).value, 1e-5f) + + it("rejects tensors with different extents along the axis"): + val tooShort = xs.slice(Axis[Time].at(0 until 3)) + the[IllegalArgumentException] thrownBy scan(Axis[Time])(zeros, (xs, tooShort))((c, x) => (c + x._1 + x._2, ())) should have message + "requirement failed: All tensors scanned over Axis[Time] must have the same extent along it, but they have 4, 3" + + it("rejects a body that changes the shape of the carry, naming the field"): + val init = Cache(Tensor(Shape(Axis[Time] -> 1, Axis[Feature] -> 3)).fill(0f), Tensor0(0)) + val error = the[Exception] thrownBy scan(Axis[Time])(init, xs): (cache, x) => + (Cache(concatenate(cache.keys, x.prependAxis(Axis[Time]), Axis[Time]), cache.count), ()) + error.getMessage should include( + "The body of a loop must not change the shape of the carry, but `keys` is Shape(Time -> 1, Feature -> 3) of Float32 and then Shape(Time -> 2, Feature -> 3) of Float32" + ) + + it("traces its body once, so the program does not grow with the number of steps"): + var traces = 0 + def scanned(steps: Int)(x: Tensor1[Feature, Float32]) = + val (last, _) = scan(Axis[Time])(x, Tensor(Shape(Axis[Time] -> steps)).fill(1.5f)): (x, factor) => + traces += 1 + ((x *! factor).sin, ()) + last + def unrolled(steps: Int)(x: Tensor1[Feature, Float32]) = (0 until steps).foldLeft(x)((x, _) => (x *! 1.5f).sin) + + jaxprSize(scanned(100)) shouldBe jaxprSize(scanned(10)) + jaxprSize(unrolled(100)) should be > jaxprSize(unrolled(10)) + traces = 0 + jit(scanned(100))(zeros) + traces shouldBe 1 + + describe("foriLoop"): + + it("fills a buffer row by row at the traced index"): + val buffer = jit((xs: Tensor2[Time, Feature, Float32]) => + foriLoop(0, 4)(Tensor(xs.shape).fill(0f)): (i, buffer) => + buffer.set(Axis[Time].at(i, 1))(xs.slice(Axis[Time].at(i, 1)) *! 2f) + )(xs) + buffer should approxEqual(xs *! 2f) + + it("is differentiable"): + def cube(x: Tensor0[Float32]): Tensor0[Float32] = foriLoop(0, 2)(x)((_, y) => y * x) + Autodiff.grad(cube)(Tensor0(2f)).value.item shouldBe 12f + + describe("whileLoop"): + + it("stops at the first carry for which the condition does not hold"): + val (steps, value) = whileLoop((Tensor0(0), Tensor0(1f)))((_, value) => value < Tensor0(100f)): + case (steps, value) => (steps + Tensor0(1), value * 2f) + steps.item shouldBe 7 + value.item shouldBe 128f + + it("stops at an iteration that depends on traced values"): + val stepsUntil = jit((limit: Tensor0[Float32]) => + whileLoop((Tensor0(0), Tensor0(1f)))((_, value) => value < limit)((steps, value) => (steps + Tensor0(1), value * 2f))._1 + ) + stepsUntil(Tensor0(100f)).item shouldBe 7 + stepsUntil(Tensor0(1000f)).item shouldBe 10 + + it("cannot be reverse-differentiated"): + def doubled(x: Tensor0[Float32]): Tensor0[Float32] = whileLoop(x)(_ < Tensor0(100f))(_ * 2f) + val error = the[Exception] thrownBy Autodiff.grad(doubled)(Tensor0(1f)) + error.getMessage should include("Reverse-mode differentiation does not work for lax.while_loop") + + describe("cond"): + + it("evaluates the branch selected by a traced predicate"): + val absolute = jit((x: Tensor0[Float32]) => cond(x > Tensor0(0f))(x)(-x)) + absolute(Tensor0(3f)).item shouldBe 3f + absolute(Tensor0(-2f)).item shouldBe 2f + + it("rejects branches with different shapes"): + val error = the[Exception] thrownBy cond(Tensor0(true))(zeros)(Tensor1(Axis[Feature]).fromArray(Array(0f))) + error.getMessage should include("The branches of cond must return the same shapes, but the result is Shape(Feature -> 3) of Float32 and then Shape(Feature -> 1) of Float32") diff --git a/examples/src/main/scala/dimwit/basic/SIRCalibration.scala b/examples/src/main/scala/dimwit/basic/SIRCalibration.scala new file mode 100644 index 00000000..00e838dd --- /dev/null +++ b/examples/src/main/scala/dimwit/basic/SIRCalibration.scala @@ -0,0 +1,81 @@ +package dimwit.examples.basic + +import dimwit.* +import dimwit.autodiff.* +import dimwit.examples.basic.SIRSimulation.* +import dimwit.optimizer.Adam +import dimwit.random.Random +import dimwit.stats.Normal + +/** Calibrates the SIR model of [[SIRSimulation]] to an observed epidemic. + * + * It finds the infection and recovery rates for which the simulated number of infectious + * individuals matches the observed one. Since the simulation is a `scan`, it can be + * differentiated, so the rates are fitted by gradient descent through the whole simulation. + */ +object SIRCalibration: + + /** The unknown parameters, on a log scale so that the rates stay positive + * + * @param logInfectionScale The log of the factor by which all infection rates are scaled + * @param logGamma The log of the recovery rate + */ + case class Params(logInfectionScale: Tensor0[Float32], logGamma: Tensor0[Float32]) + + /** The number of infectious individuals after each time step, simulated with the parameters */ + def infectiousOverTime( + initial: SIRState, + beta: Tensor2[Group, Prime[Group], Float32], + dt: Tensor0[Float32], + contact: Tensor1[Time, Float32] + )(params: Params): Tensor1[Time, Float32] = + val (_, infectious) = simulate(initial, beta *! params.logInfectionScale.exp, params.logGamma.exp, dt, contact) + infectious.sum(Axis[Group]) + + /** The mean squared difference between the simulated and the observed fraction of infectious individuals */ + def loss( + initial: SIRState, + beta: Tensor2[Group, Prime[Group], Float32], + dt: Tensor0[Float32], + contact: Tensor1[Time, Float32], + observed: Tensor1[Time, Float32] + )(params: Params): Tensor0[Float32] = + val population = initial.N.sum + ((infectiousOverTime(initial, beta, dt, contact)(params) - observed) /! population).pow(2f).mean + + @main def runSIRCalibration(): Unit = + dimwit.initialize() + + val initial = exampleInitialState() + val beta = exampleBeta() + val dt = Tensor0(0.1f) + val nSteps = 300 + val contact = Tensor(Shape(Axis[Time] -> nSteps)).fill(1f) + + // Observe an epidemic with known rates, with some noise + + val trueParams = Params(logInfectionScale = Tensor0(0f), logGamma = Tensor0(0.1f).log) + val noise = Normal.standardNormal(Shape(Axis[Time] -> nSteps)).sample(Random.Key(0)) *! 20f + val observed = infectiousOverTime(initial, beta, dt, contact)(trueParams) + noise + + // Fit the rates, starting from a wrong guess + + val initialGuess = Params(logInfectionScale = Tensor0(0.5f).log, logGamma = Tensor0(0.2f).log) + val trainLoss = jit(loss(initial, beta, dt, contact, observed)) + val adam = Adam(learningRate = 0.05f) + + def report(params: Params): String = + f"infection scale = ${params.logInfectionScale.exp.item}%.3f, gamma = ${params.logGamma.exp.item}%.3f" + + val fitted = adam + .iterate(initialGuess)(grad(trainLoss)) + .zipWithIndex + .tapEach: + case (params, index) => + if index % 25 == 0 then println(f"step $index%3d: loss = ${trainLoss(params).item}%.2e, ${report(params)}") + .map(_._1) + .drop(150) + .next() + + println(s"fitted: ${report(fitted)}") + println(s"true: ${report(trueParams)}") diff --git a/examples/src/main/scala/dimwit/basic/SIRSimulation.scala b/examples/src/main/scala/dimwit/basic/SIRSimulation.scala index 43e6d0cc..5fc2331f 100644 --- a/examples/src/main/scala/dimwit/basic/SIRSimulation.scala +++ b/examples/src/main/scala/dimwit/basic/SIRSimulation.scala @@ -46,93 +46,96 @@ object SIRSimulation: SIRState(SNext, INext, RNext) - /** run n steps of the simulation, starting from the initial state + /** run the simulation over the time steps of `contact`, starting from the initial state + * + * The loop is a `scan`: its body is traced once, so the simulation can be jitted and + * differentiated, however many steps it runs. * * @param initial The initial state of the system * @param beta @see [[step]] * @param gamma @see [[step]] * @param dt @see [[step]] - * @param nSteps The number of steps to simulate - * @return The trajectory of the system over time + * @param contact The factor by which contacts, and thus infections, are scaled at each time step + * (e.g. 1 without and below 1 during an intervention) + * @return The final state of the system, and the infectious individuals per group after each time step */ def simulate( initial: SIRState, beta: Tensor2[Group, Prime[Group], Float32], gamma: Tensor0[Float32], - dt: Tensor0[Float32] - ): LazyList[SIRState] = - LazyList.iterate(initial)(step(beta, gamma, dt)) - - @main def runSIRSimulation(): Unit = - dimwit.initialize() - - val groupDim = Axis[Group] -> 3 - - // Setup initial state of the system - - val initialS = Tensor(Shape(groupDim)).fromFunction(index => + dt: Tensor0[Float32], + contact: Tensor1[Time, Float32] + ): (SIRState, Tensor2[Time, Group, Float32]) = + scan(Axis[Time])(initial, contact): (state, contactFactor) => + val next = step(beta *! contactFactor, gamma, dt)(state) + (next, next.I) + + /** The extent of the group axis in the example: children, adults and elderly */ + val groupExtent = Axis[Group] -> 3 + + /** The initial state of the example: a few infectious children and adults */ + def exampleInitialState(): SIRState = + val initialS = Tensor(Shape(groupExtent)).fromFunction(index => index(Axis[Group]) match case 0 => 990f case 1 => 1995f case 2 => 1500f ) - val initialI = Tensor(Shape(groupDim)).fromFunction(index => + val initialI = Tensor(Shape(groupExtent)).fromFunction(index => index(Axis[Group]) match case 0 => 10f case 1 => 5f case 2 => 0f ) - val initialR = Tensor(Shape(groupDim)).fromFunction(_ => 0f) - val initial = SIRState(initialS, initialI, initialR) - - // Setup parameters for the simulation - - /* - * beta(h, g) controls how strongly infectious individuals in group h - * infect susceptible individuals in group g. - */ - val beta: Tensor2[Group, Prime[Group], Float32] = - Tensor(Shape(groupDim, Axis[Prime[Group]] -> groupDim.size)).fromFunction(index => - (index(Axis[Group]), index(Axis[Prime[Group]])) match - // infectious children -> susceptible children/adults/elderly - case (0, 0) => 0.40f - case (0, 1) => 0.20f - case (0, 2) => 0.10f - - // infectious adults -> susceptible children/adults/elderly - case (1, 0) => 0.20f - case (1, 1) => 0.30f - case (1, 2) => 0.15f - - // infectious elderly -> susceptible children/adults/elderly - case (2, 0) => 0.10f - case (2, 1) => 0.15f - case (2, 2) => 0.20f - - case (_, _) => throw new IllegalArgumentException("Invalid group indices") - ) + val initialR = Tensor(Shape(groupExtent)).fill(0f) + SIRState(initialS, initialI, initialR) + + /** The infection rates of the example + * + * beta(h, g) controls how strongly infectious individuals in group h + * infect susceptible individuals in group g. + */ + def exampleBeta(): Tensor2[Group, Prime[Group], Float32] = + Tensor(Shape(groupExtent, Axis[Prime[Group]] -> groupExtent.size)).fromFunction(index => + (index(Axis[Group]), index(Axis[Prime[Group]])) match + // infectious children -> susceptible children/adults/elderly + case (0, 0) => 0.40f + case (0, 1) => 0.20f + case (0, 2) => 0.10f + + // infectious adults -> susceptible children/adults/elderly + case (1, 0) => 0.20f + case (1, 1) => 0.30f + case (1, 2) => 0.15f + + // infectious elderly -> susceptible children/adults/elderly + case (2, 0) => 0.10f + case (2, 1) => 0.15f + case (2, 2) => 0.20f + + case (_, _) => throw new IllegalArgumentException("Invalid group indices") + ) + + @main def runSIRSimulation(): Unit = + dimwit.initialize() + + val initial = exampleInitialState() + val beta = exampleBeta() val gamma = Tensor0(0.1f) val dt = Tensor0(0.1f) - // Run the simulation + // Run the simulation, without and with an intervention that reduces contacts for a while - val nSteps = 160 - val trajectory = - SIRSimulation.simulate( - initial = initial, - beta = beta, - gamma = gamma, - dt = dt - ).take(nSteps + 1).toList + val nSteps = 1000 + val noIntervention = Tensor(Shape(Axis[Time] -> nSteps)).fill(1f) + val lockdown = Tensor1(Axis[Time]).fromArray(Array.tabulate(nSteps)(t => if t >= 100 && t < 300 then 0.4f else 1f)) - // Report the results + for (name, contact) <- List("No intervention" -> noIntervention, "Lockdown" -> lockdown) do + val (finalState, infectious) = simulate(initial, beta, gamma, dt, contact) - val infectedOverTime: Tensor1[Time, Float32] = - stack(trajectory.map(_.I), Axis[Time]).sum(Axis[Group]) - println(s"I(0) = ${infectedOverTime(Axis[Time].at(0))}") - println(s"I(mid) = ${infectedOverTime(Axis[Time].at(nSteps / 2))}") - println(s"I(end) = ${infectedOverTime(Axis[Time].at(nSteps))}") + // Report the results - val finalInfectedByGroup: Tensor1[Group, Float32] = - stack(trajectory.map(_.I), Axis[Time]).select(Axis[Time].at(nSteps)) - println(s"Final infected by group: $finalInfectedByGroup") + val infectedOverTime: Tensor1[Time, Float32] = infectious.sum(Axis[Group]) + println(s"$name:") + println(s" peak of infectious = ${infectedOverTime.max} at time step ${infectedOverTime.argmax(Axis[Time])}") + println(s" recovered at the end by group = ${finalState.R}") diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index ee51ba84..c4508d31 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -660,6 +660,68 @@ println(s"Reduced A: ${summedA}") println(s"Reduced B: ${summedB}") ``` +### Loops and Branches: scan, foriLoop, whileLoop, cond + +A Scala loop inside `jit` is unrolled while tracing: 100 iterations become 100 copies of the body, +and the compile time grows with them. `scan`, `foriLoop`, `whileLoop` and `cond` wrap the `jax.lax` +primitives of the same names, whose body is traced once. They work eagerly and under `jit`. + +The state carried from one iteration to the next may be any tensor tree (a tensor, a tuple, a case class). +Since the body is traced once, it must not change the shape of the carry. + +```scala mdoc:reset:silent +import dimwit.* + +trait Time derives Label +trait Feature derives Label + +val signal = Tensor1(Axis[Time]).fromArray(Array(1.0f, 3.0f, 2.0f, 6.0f)) + +// scan: loop over the axis Time of a tensor. The body gets the carry and the slice at each +// position along Time, and returns the next carry and an output; the outputs are stacked along Time. +val (lastAverage, averages) = scan(Axis[Time])(Tensor0(0.0f), signal): (average, x) => + val next = average * 0.5f + x * 0.5f + (next, next) +// averages: Tensor1[Time, Float32] = [0.5, 1.75, 1.875, 3.9375] + +// scan over several tensors along the same axis, like zipvmap; () means no outputs +val weights = Tensor1(Axis[Time]).fromArray(Array(1.0f, 1.0f, 2.0f, 2.0f)) +val (weightedSum, _) = scan(Axis[Time])(Tensor0(0.0f), (signal, weights)): + case (sum, (x, weight)) => (sum + x * weight, ()) + +// foriLoop: loop over the indices lower until upper. The index is a traced Tensor0[Int32], +// e.g. to write a window at it +val squares = foriLoop(0, 4)(Tensor(Shape(Axis[Time] -> 4)).fill(0)): (i, buffer) => + buffer.set(Axis[Time].at(i, 1))((i * i).broadcastTo(Shape(Axis[Time] -> 1))) +// squares: [0, 1, 4, 9] + +// whileLoop: loop while a condition holds, here Newton's method for the square root of 2 +val (root, iterations) = whileLoop((Tensor0(1.0f), Tensor0(0)))((x, _) => (x * x - Tensor0(2.0f)).abs > Tensor0(1e-5f)): + case (x, n) => ((x + Tensor0(2.0f) / x) * 0.5f, n + Tensor0(1)) + +// cond: choose a branch by a predicate that may be traced, unlike Scala's if; +// both branches must return the same shapes +val relu = jit((x: Tensor0[Float32]) => cond(x > Tensor0(0.0f))(x)(Tensor0(0.0f))) +``` + +`scan` and `foriLoop` are differentiable, so `Autodiff.grad` through them gives the same gradient as through the +unrolled loop. JAX cannot reverse-differentiate `whileLoop`, since the number of iterations is only known at runtime; +use forward mode (`Autodiff.jacFwd`) for it instead. + +```scala mdoc:fail +// ERROR: the scanned axis must exist in the tensor +val notAnAxis = scan(Axis[Feature])(Tensor0(0.0f), signal)((sum, x) => (sum + x, ())) +``` + +```scala mdoc +// ERROR at runtime: the body must not change the shape of the carry. +// (The error is raised inside the traced body; here only its message is shown.) +val shapeError = scala.util.Try: + scan(Axis[Time])(Tensor1(Axis[Feature]).fromArray(Array(0.0f)), signal): (state, x) => + (concatenate(state, state, Axis[Feature]), ()) +.failed.get.getMessage.linesIterator.map(_.trim).find(_.contains("requirement failed")).get +``` + --- ## Automatic Differentiation From 92dd7cde668bf2c46752861fde03908de4ab0019 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Wed, 7 Oct 2026 19:19:58 +0200 Subject: [PATCH 2/2] Add structured control flow: scan, foriLoop, whileLoop and cond --- AGENTS.md | 2 +- .../tensor/TensorOpsControlFlowSuite.scala | 6 +- mdocs/AGENTS.md | 2 +- uv.lock | 123 ++++++++++++++++++ 4 files changed, 128 insertions(+), 5 deletions(-) create mode 100644 uv.lock diff --git a/AGENTS.md b/AGENTS.md index b5ef95f2..a5cb9ada 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -925,7 +925,7 @@ val (weightedSum, _) = scan(Axis[Time])(Tensor0(0.0f), (signal, weights)): // foriLoop: loop over the indices lower until upper. The index is a traced Tensor0[Int32], // e.g. to write a window at it val squares = foriLoop(0, 4)(Tensor(Shape(Axis[Time] -> 4)).fill(0)): (i, buffer) => - buffer.set(Axis[Time].at(i, 1))((i * i).broadcastTo(Shape(Axis[Time] -> 1))) + buffer.updated(Axis[Time].at(i, 1))((i * i).broadcastTo(Shape(Axis[Time] -> 1))) // squares: [0, 1, 4, 9] // whileLoop: loop while a condition holds, here Newton's method for the square root of 2 diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsControlFlowSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsControlFlowSuite.scala index a21857f9..1ab0af14 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsControlFlowSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsControlFlowSuite.scala @@ -59,7 +59,7 @@ class TensorOpsControlFlowSuite extends DimwitTest: (Cache(keys, cache.count + Tensor0(1)), cache.count) val weightedSum = xs.unstack(Axis[Time]).zip(Seq(1f, 2f, 3f, 4f)).map((x, w) => x *! w).reduce(_ + _) cache.count.item shouldBe 4 - cache.keys.slice(Axis[Time].at(0)) should approxEqual(weightedSum) + cache.keys(Axis[Time].at(0)) should approxEqual(weightedSum) counts shouldEqual Tensor1(Axis[Time]).fromArray(Array(0, 1, 2, 3)) it("has the same gradient as the unrolled loop"): @@ -72,7 +72,7 @@ class TensorOpsControlFlowSuite extends DimwitTest: Autodiff.grad(scanned)(w).value should approxEqual(Autodiff.grad(unrolled)(w).value, 1e-5f) it("rejects tensors with different extents along the axis"): - val tooShort = xs.slice(Axis[Time].at(0 until 3)) + val tooShort = xs(Axis[Time].at(0 until 3)) the[IllegalArgumentException] thrownBy scan(Axis[Time])(zeros, (xs, tooShort))((c, x) => (c + x._1 + x._2, ())) should have message "requirement failed: All tensors scanned over Axis[Time] must have the same extent along it, but they have 4, 3" @@ -104,7 +104,7 @@ class TensorOpsControlFlowSuite extends DimwitTest: it("fills a buffer row by row at the traced index"): val buffer = jit((xs: Tensor2[Time, Feature, Float32]) => foriLoop(0, 4)(Tensor(xs.shape).fill(0f)): (i, buffer) => - buffer.set(Axis[Time].at(i, 1))(xs.slice(Axis[Time].at(i, 1)) *! 2f) + buffer.updated(Axis[Time].at(i, 1))(xs(Axis[Time].at(i, 1)) *! 2f) )(xs) buffer should approxEqual(xs *! 2f) diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index c4508d31..91e190fb 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -692,7 +692,7 @@ val (weightedSum, _) = scan(Axis[Time])(Tensor0(0.0f), (signal, weights)): // foriLoop: loop over the indices lower until upper. The index is a traced Tensor0[Int32], // e.g. to write a window at it val squares = foriLoop(0, 4)(Tensor(Shape(Axis[Time] -> 4)).fill(0)): (i, buffer) => - buffer.set(Axis[Time].at(i, 1))((i * i).broadcastTo(Shape(Axis[Time] -> 1))) + buffer.updated(Axis[Time].at(i, 1))((i * i).broadcastTo(Shape(Axis[Time] -> 1))) // squares: [0, 1, 4, 9] // whileLoop: loop while a condition holds, here Newton's method for the square root of 2 diff --git a/uv.lock b/uv.lock new file mode 100644 index 00000000..559c2fde --- /dev/null +++ b/uv.lock @@ -0,0 +1,123 @@ +version = 1 +requires-python = "==3.13.*" + +[[package]] +name = "dimwit-python-env" +version = "0.1.0" +source = { virtual = "." } +dependencies = [ + { name = "einops" }, + { name = "jax" }, +] + +[package.metadata] +requires-dist = [ + { name = "einops", specifier = ">=0.8.1" }, + { name = "jax", specifier = ">=0.8.2" }, +] + +[[package]] +name = "einops" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638 }, +] + +[[package]] +name = "jax" +version = "0.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jaxlib" }, + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "opt-einsum" }, + { name = "scipy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/90/aa/c3b6ac3863a490661e255d2356802b54d57b533460fa221d30743d81a6c1/jax-0.11.1.tar.gz", hash = "sha256:43e0d45b1dac002eca04c2de73298395225dd0d4651e3f573a540eaa18abfd80", size = 2864119 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/f4/5f9a286e1eeca9faa8b83885edccf63ab2f1b22e17f6a5ea66ac86b1afe4/jax-0.11.1-py3-none-any.whl", hash = "sha256:72ed60bf85a0b53d0f5b5cd61e37a8b25f369f6f88481feb98ab489894861a82", size = 3302513 }, +] + +[[package]] +name = "jaxlib" +version = "0.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "scipy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/22/d076bf545853d1f60ef906cb6d69751f9b149f857f4390de989d5165ed80/jaxlib-0.11.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c7518af3168ab4cf0a00b92f5463018baf0dad3faa3b358a3d688996afa6fd50", size = 63174384 }, + { url = "https://files.pythonhosted.org/packages/a8/3c/5b0ce63c880d4f9cef492f04e11be7a60bdd3f529311839652b928e8b320/jaxlib-0.11.1-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:3a679a49b8d329b6e431f3aeb666467c6efac750e3b2ee90d0e74efe7959e643", size = 82774089 }, + { url = "https://files.pythonhosted.org/packages/f5/28/ef1371618e784880e68c2ab5277d5ccbaf518a2a5ec9ec05ef90a9ee7d1b/jaxlib-0.11.1-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:7ebe86a7b891fcda83e7f634a677d285681a7b80b3ec9ed435536078a8371acf", size = 87901562 }, + { url = "https://files.pythonhosted.org/packages/fe/c7/f92d23dc89bdbe5193ad139aa8973eab485898f330de2033bc4fdd8ac4c9/jaxlib-0.11.1-cp313-cp313-win_amd64.whl", hash = "sha256:8379a8902db216f21673e7591acedc8f925658c59d9d2ec9afd5e2244246986a", size = 68531205 }, +] + +[[package]] +name = "ml-dtypes" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/72/307d7c4bd0600601c7133fba5cb78af7db968152951c1cd473abb1cda782/ml_dtypes-0.6.0.tar.gz", hash = "sha256:5e60251d32ced5598972e4d5e06a2f044341f9291402551a3f6f0ec44f9299b0", size = 3032327 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/50/51/fd1582b8f5ed8a9e7be0e161a6ea0dff70cb280479a12178df0b3a72700e/ml_dtypes-0.6.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:084dfe51a7ad58b171f05115f8226ed4233a454a1611371947e806e76f0c638d", size = 565468 }, + { url = "https://files.pythonhosted.org/packages/d2/22/20fd70ca6ed12446cb92d5b2a7745bd185f9d8b8cdeeadad976574398e6b/ml_dtypes-0.6.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:28d676428b104bb9717b0928bc5c5129f2d6b51b6727587cc4289e7bf8713cb5", size = 360232 }, + { url = "https://files.pythonhosted.org/packages/89/a5/da8ae6c6f1babe4b68e3e55d43d39b529e29774f10e0910671a6b8c86eb8/ml_dtypes-0.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:26b1f1fa4f0435a2946859823f6e2bf06796f1e9f10f5a05b08a5e3c8f46ff69", size = 410169 }, + { url = "https://files.pythonhosted.org/packages/e2/55/4561acefa00fa4bcbfb82ca6a48578b41f372cd7dd7cdd6eb4720abc2e5f/ml_dtypes-0.6.0-cp313-cp313-win_amd64.whl", hash = "sha256:fb87f46b4f7ad7b5d3ad8f4b452b024bd4229d44c8ff934798c1fe656210387a", size = 439357 }, + { url = "https://files.pythonhosted.org/packages/b1/5d/6a01538e507ef0ed5e879985b13a92467bf8960696fb1131f8b8cadc60ff/ml_dtypes-0.6.0-cp313-cp313-win_arm64.whl", hash = "sha256:57ed0d6b4ac5e7868361303a9c57fbcf63b768236ee14456f585dfcf260d0292", size = 552278 }, +] + +[[package]] +name = "numpy" +version = "2.5.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/13/01/11703282db468b85f6f7b8c7f22d058de5970d5c7e60a3a8aaa313c3de36/numpy-2.5.3.tar.gz", hash = "sha256:df2d5874ff183595a4ba404edd04f6bd9b5505c1d7708573f6a6c17489a67563", size = 20791231 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/79/e5/8fb89cd46d14e35699d13bf943a5f5f441ecee8667120a1f6105ab89e349/numpy-2.5.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:66a78fe4556c60aceda5916f9eacd638b18e9e681016ec302dcb4682d6d4d034", size = 16991061 }, + { url = "https://files.pythonhosted.org/packages/2f/06/9dc9e48b5e5e941c8b10350c5ff2d721da42a20517d911d15544246775ff/numpy-2.5.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:92f30e89b8ee0ecf363033576c422b2f58fed6a80bed0aa48dff6d14c654663e", size = 12003676 }, + { url = "https://files.pythonhosted.org/packages/ab/2a/98282aa5b8f58b1157d440bb6282eed47e3632a5de53a714fbab17e659fe/numpy-2.5.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f9a2353b37a1a9e78fd82b27ad7e2a32a2d036604d18f02b05e3136c62ca3b09", size = 5439695 }, + { url = "https://files.pythonhosted.org/packages/a1/f9/b6533d777be9d6ffd29dc1be0867e563e6e8cc9a220ff1b716adc317f060/numpy-2.5.3-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:ccbc4665079665c3cf3bab4db9f6b095370cd6437d66be549b6c2a1fd19e1958", size = 6779395 }, + { url = "https://files.pythonhosted.org/packages/73/85/735720d04ec197c5dcfacdfc9922667c7f1f5f496a279b7ba4d7c74c4cc7/numpy-2.5.3-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c76d5dde9f445058f83d0c02af00557a4db91de9a9a57c0df87d1535001d654b", size = 15681750 }, + { url = "https://files.pythonhosted.org/packages/3a/1b/3b16a9bc514a440a7a0883684111dcb1ef1aee960af2ca95da8fc775f124/numpy-2.5.3-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a5fa86b80fd24bcd1aff83ad23be44ea323de3f787be8f8b15d4a65621e25321", size = 16708577 }, + { url = "https://files.pythonhosted.org/packages/69/c4/386f397831b07328b639c96c5b62719346cf4baf07c68d927239752b1534/numpy-2.5.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:bd4cb9ad3c7889b9b3fe0a9a9fb5d2ed26f9879bff2608d9f01aed147a20d231", size = 17042047 }, + { url = "https://files.pythonhosted.org/packages/5f/3e/a700ecbf36e85ae8328fd3b0e12eeddc22ed6358a64cb2bd913e0d195d65/numpy-2.5.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1302b90c0e52281681b2975adfe8a860cb7b12216a27b4b0b4207c44bf7bccf0", size = 18465724 }, + { url = "https://files.pythonhosted.org/packages/41/ee/38e785e88a4045f6ad1d1f2808dcdfafdca48c760260c0587bf171e29fc9/numpy-2.5.3-cp313-cp313-win32.whl", hash = "sha256:1c80eabb4035ecf4ca9cd49cde8a9fdd69a729e63e6474887d1523ade7aa277f", size = 6129003 }, + { url = "https://files.pythonhosted.org/packages/f3/ec/100f2b1794ede74a9b3d7ec6b9736927f56713414c1dfe19ab6c383494bf/numpy-2.5.3-cp313-cp313-win_amd64.whl", hash = "sha256:71cad2b2a7451ab79d8f5e71b453485b6775963d5cf794179144a7463fe6e8ec", size = 12560965 }, + { url = "https://files.pythonhosted.org/packages/80/b1/7dc825ca94c12acebbce4c37caa5e198695eb31424bc579679f32b1bb49d/numpy-2.5.3-cp313-cp313-win_arm64.whl", hash = "sha256:8e4dd766076855b5ff7ea52fa5f07ce26286726e0f8bff446b7739d02e6ea204", size = 10482343 }, +] + +[[package]] +name = "opt-einsum" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/b9/2ac072041e899a52f20cf9510850ff58295003aa75525e58343591b0cbfb/opt_einsum-3.4.0.tar.gz", hash = "sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac", size = 63004 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd", size = 71932 }, +] + +[[package]] +name = "scipy" +version = "1.18.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7e/74/66de6258867beb2ef08f35f9f2ac017a52cacd5081714d239ff1a442d458/scipy-1.18.1.tar.gz", hash = "sha256:52c4b7422442aba924d03ad4019852b08a92e64ea187b933135687bfe2747307", size = 30781235 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b6/55/4540ee0f9c42a9ad7109d0d1a8cc70de54c3572b01c6693a2b1c70e90ceb/scipy-1.18.1-cp313-cp313-macosx_10_15_x86_64.whl", hash = "sha256:3ab3523da44749156e1f68b464dc56af11ae4cbc5c739a49d05f32b982eca9f3", size = 31089958 }, + { url = "https://files.pythonhosted.org/packages/2a/f5/769f36d14922b8071a43e95d24d18b6bdafad10d7f5cf647867e1ac052bc/scipy-1.18.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:e6fb6a55cc0ba97b59a1f288fb86dc6fce8bdfc0fffcbfd015e3a954bf2a2d93", size = 28715106 }, + { url = "https://files.pythonhosted.org/packages/9a/d7/21d890274f75ea37a8209d5519e72da3da90302e3b9fb8397a0918386a62/scipy-1.18.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:ea324d9dd34c38bfb9bec8ca4d1b407db97dbb74029f566b8e322b1b6fe56fe6", size = 20456846 }, + { url = "https://files.pythonhosted.org/packages/ec/01/798430ecea2e78ec7c02663d5f71c007bb6abeca931080debd40d7fa55ea/scipy-1.18.1-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:75b00eb8fb802090aa903f4ea1c7f5a584779f967361e68b7e98e531cc2d7174", size = 23087986 }, + { url = "https://files.pythonhosted.org/packages/e6/5f/4634e9d35c68496e4e34cb6946eafab044458e6cedab42b40b6588e475b6/scipy-1.18.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d416b16cccfd70fbf62400e84d0bb2f4e6af519a45557f1692c749b37f14b315", size = 33998146 }, + { url = "https://files.pythonhosted.org/packages/41/48/6450ed9243315322bbc19ac57b9b70d66a20bf1d38d124c96bc4bf6af9ea/scipy-1.18.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fdaf5ea890a6183d0565f51a61799d67081bd5b1cf03c5f4b3fd3732108625c9", size = 35312578 }, + { url = "https://files.pythonhosted.org/packages/00/bd/bf5a4be6a3525676499f6dff307991739ff6fdcad1481b1aeb6745339f58/scipy-1.18.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:c825cef2f49e46753726a7181a8e199804a912b29519ada542c6ebc654951899", size = 35612621 }, + { url = "https://files.pythonhosted.org/packages/bd/4e/3c45c33e00a77996c4b1cb707929f833ba7b1d522ee29f882512c330676d/scipy-1.18.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e3b417bf8c2c7c16e8f58ad91db17783ec911ac16e7b50eb6eab6e809b4f5b07", size = 37457323 }, + { url = "https://files.pythonhosted.org/packages/93/0e/e0348fbc0dbab65c114cf78957e7dfeb49f8e8b556b4d930cc12ff195e18/scipy-1.18.1-cp313-cp313-win_amd64.whl", hash = "sha256:559ed65f60c1af5a03f3912605a1b5114f522c7c32fb23c3376ae8f03219fe28", size = 36622841 }, + { url = "https://files.pythonhosted.org/packages/50/a8/6a77f5f267c555108f0a864b6db714363dab567a8266422a79a385f9232b/scipy-1.18.1-cp313-cp313-win_arm64.whl", hash = "sha256:cd479fc04dd9401e3b4f49e76518768ef99c4f517a98c284eb091fd725719adf", size = 24399315 }, +]