From 281b11ce7562f8eba1cb801efd19e37efe1d5512 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Thu, 13 Aug 2026 13:01:20 +0200 Subject: [PATCH] Support NamedTuples as return types in vmap and zipvmap Replaced the `PrependAxis` match type with a typeclass to overcome resolution limitations and explicitly support Scala 3 NamedTuples. Functions passed to `vmap` and `zipvmap` can now return named tuples (e.g., `(first = x, second = y)`), and the operation will correctly prepend the mapped axis to all enclosed tensors. Added test coverage to verify the new behavior. --- .../tensor/tensorops/FunctionalOps.scala | 61 +++++++++++++------ .../tensor/TensorOpsFunctionalSuite.scala | 10 ++- 2 files changed, 52 insertions(+), 19 deletions(-) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala index d882011..36dcb2f 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala @@ -11,25 +11,44 @@ import dimwit.tensor.ShapeTypeHelpers.AxisReplacer import dimwit.tensor.ShapeTypeHelpers.SharedAxisRemover import dimwit.tensor.Tensor import dimwit.tensor.Tensor0 -import dimwit.tensor.tensorops.FunctionalOps.ZipVmap.TensorsOf +import dimwit.tensortree.TensorTree import me.shadaj.scalapy.py import me.shadaj.scalapy.py.SeqConverters import me.shadaj.scalapy.readwrite.Reader import me.shadaj.scalapy.readwrite.Writer -import dimwit.tensortree.TensorTree -import dimwit.tensor.ShapeTypeHelpers.UnwrapAxes -import dimwit.tensor.ShapeTypeHelpers.AxesRemover + +import scala.NamedTuple.NamedTuple object FunctionalOps: - type PrependAxes[Axes <: Tuple, FOut] = Axes match - case EmptyTuple => FOut - case h *: t => PrependAxis[h, PrependAxes[t, FOut]] + trait PrependAxis[L, FOut]: + type Out + + object PrependAxis: + type Aux[L, FOut, Out0] = PrependAxis[L, FOut] { type Out = Out0 } + + // Case 1: Tensor + given tensorCase[L, Shape <: Tuple, V]: PrependAxis[L, Tensor[Shape, V]] with + type Out = Tensor[L *: Shape, V] - type PrependAxis[L, FOut] = FOut match - case Tensor[shape, v] => Tensor[L *: shape, v] - case EmptyTuple => EmptyTuple - case h *: t => PrependAxis[L, h] *: PrependAxis[L, t] + // Case 2: EmptyTuple + given emptyTupleCase[L]: PrependAxis[L, EmptyTuple] with + type Out = EmptyTuple + + // Case 3: Recursive Tuple (simplified without explicit tpOut variable) + given tupleCase[L, H, T <: Tuple, tpOut <: Tuple](using + hp: PrependAxis[L, H], + tp: PrependAxis.Aux[L, T, tpOut] + ): PrependAxis.Aux[L, H *: T, hp.Out *: tpOut] = + new PrependAxis[L, H *: T]: + type Out = hp.Out *: tpOut + + // Case 4: NamedTuple + given namedTupleCase[L, Names <: Tuple, Values <: Tuple, vpOut <: Tuple](using + vp: PrependAxis.Aux[L, Values, vpOut] + ): PrependAxis.Aux[L, NamedTuple[Names, Values], NamedTuple[Names, vpOut]] = + new PrependAxis[L, NamedTuple[Names, Values]]: + type Out = NamedTuple[Names, vpOut] object ZipVmap: @@ -71,10 +90,12 @@ object FunctionalOps: ev: SharedAxisRemover[ShapesOf[Inputs], L] )( f: TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]] => FOut + )(using + prependAxis: PrependAxis[L, FOut] )(using toPyTree: TensorTree[FOut], - fromPyTree: TensorTree[PrependAxis[L, FOut]] - ): PrependAxis[L, FOut] = + fromPyTree: TensorTree[prependAxis.Out] + ): prependAxis.Out = val fpy = (args: py.Dynamic) => OnError.traceStack: val tensorList = args.as[Seq[py.Dynamic]].zip(ev.shapesLabels).map: (jaxArr, labels) => @@ -110,11 +131,13 @@ object FunctionalOps: )(using ev: SharedAxisRemover[(T, T2), L] )( - f: TensorsOf[ev.RemainingAxes, (V, V)] => FOut + f: ZipVmap.TensorsOf[ev.RemainingAxes, (V, V)] => FOut + )(using + prependAxis: PrependAxis[L, FOut] )(using toPyTree: TensorTree[FOut], - fromPyTree: TensorTree[PrependAxis[L, FOut]] - ): PrependAxis[L, FOut] = + fromPyTree: TensorTree[prependAxis.Out] + ): prependAxis.Out = ZipVmap.zipvmap(axis)(t, other)(f) /** Vectorized mapping over a specified axis of the tensor. @@ -129,11 +152,13 @@ object FunctionalOps: ev: AxisRemover[T, VmapAxis] )( f: Tensor[ev.RemainingAxes, V] => FOut + )(using + prependAxis: PrependAxis[VmapAxis, FOut] )(using toPyTree: TensorTree[FOut], - fromPyTree: TensorTree[PrependAxis[VmapAxis, FOut]], + fromPyTree: TensorTree[prependAxis.Out], labels: Labels[ev.RemainingAxes] - ): PrependAxis[VmapAxis, FOut] = + ): prependAxis.Out = val fpy = (jxpr: Jax.PyDynamic) => OnError.traceStack: val innerTensor = Tensor[ev.RemainingAxes, V](jxpr) diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala index 0278b1f..59d14a0 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala @@ -25,10 +25,18 @@ class TensorOpsFunctionalSuite extends DimwitTest: it("vmap return tuple"): val t = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(0f) val (y1, y2) = t.vmap(Axis[A]): x => - (x +! 5f, x -! 5f) + val tt = (x +! 5f, x -! 5f) + tt y1 shouldEqual (t +! 5f) y2 shouldEqual (t -! 5f) + it("vmap return named tuple"): + val t = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(0f) + val res = t.vmap(Axis[A]): x => + (first = x +! 5f, second = x -! 5f) + res.first shouldEqual (t +! 5f) + res.second shouldEqual (t -! 5f) + it("vmap over Axis B (columns)"): val res = t2.vmap(Axis[B])(_.sum) res shouldEqual Tensor1(Axis[B]).fromArray(Array(4.0f, 6.0f))