From 4c26e711b354c513b751add806583fe06a0b3a24 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Fri, 7 Aug 2026 13:58:09 +0200 Subject: [PATCH 1/3] Change zipvmap to allow returning tuples of tensors --- .../tensor/tensorops/FunctionalOps.scala | 51 ++++++++++++++++--- 1 file changed, 44 insertions(+), 7 deletions(-) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala index eec4b3d..352047d 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala @@ -17,6 +17,41 @@ import me.shadaj.scalapy.py.SeqConverters import me.shadaj.scalapy.readwrite.Reader import me.shadaj.scalapy.readwrite.Writer +trait ZipVmapResult[L: Label, FOut, MOut]: + type MappedOut + def toPy(out: FOut): py.Dynamic + def fromPy(pyOut: py.Dynamic): MOut + +object ZipVmapResult: + + // Single Tensor case + given singleTensor[L: Label, Shape <: Tuple: Labels, V]: ZipVmapResult[L, Tensor[Shape, V], Tensor[L *: Shape, V]] with + def toPy(out: Tensor[Shape, V]): py.Dynamic = out.jaxValue + def fromPy(pyOut: py.Dynamic): Tensor[L *: Shape, V] = Tensor(pyOut) + + // Empty Tuple case + given emptyTuple[L: Label]: ZipVmapResult[L, EmptyTuple, EmptyTuple] with + def toPy(out: EmptyTuple): py.Dynamic = py.Dynamic.global.tuple(Seq.empty[py.Dynamic].toPythonProxy) + def fromPy(pyOut: py.Dynamic): EmptyTuple = EmptyTuple + + // Inductive Tuple case (Pairs, Triples, N-tuples) + given consTuple[L: Label, H, HOut, T <: Tuple, TOut <: Tuple](using + hRes: ZipVmapResult[L, H, HOut], + tRes: ZipVmapResult[L, T, TOut] + ): ZipVmapResult[L, H *: T, HOut *: TOut] with + + def toPy(out: H *: T): py.Dynamic = + val headPy = hRes.toPy(out.head) + val tailSeq = tRes.toPy(out.tail).as[Seq[py.Dynamic]] + py.Dynamic.global.tuple((headPy +: tailSeq).toPythonProxy) + + def fromPy(pyOut: py.Dynamic): HOut *: TOut = + val seq = pyOut.as[Seq[py.Dynamic]] + val h = hRes.fromPy(seq.head) + val tailPy = py.Dynamic.global.tuple(seq.tail.toPythonProxy) + val t = tRes.fromPy(tailPy) + h *: t + object FunctionalOps: object ZipVmap: @@ -51,15 +86,17 @@ object FunctionalOps: * ... * } */ - def zipvmap[L: Label, Inputs <: Tuple, OutShape <: Tuple: Labels, OutV]( + def zipvmap[L: Label, Inputs <: Tuple, FOut, MOut]( axis: Axis[L] )( - tensors: Inputs // This is a Tuple of Tensors + tensors: Inputs )(using ev: SharedAxisRemover[ShapesOf[Inputs], L] )( - f: TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]] => Tensor[OutShape, OutV] - ): Tensor[L *: OutShape, OutV] = + f: TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]] => FOut + )(using + outMapper: ZipVmapResult[L, FOut, MOut] + ): MOut = val fpy = (args: py.Dynamic) => OnError.traceStack: val tensorList = args.as[Seq[py.Dynamic]].zip(ev.shapesLabels).map: (jaxArr, labels) => @@ -67,17 +104,17 @@ object FunctionalOps: val inputTuple = Tuple.fromArray(tensorList.toArray) val result = f(inputTuple.asInstanceOf[TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]]]) - result.jaxValue + outMapper.toPy(result) val jaxInputs = py.Dynamic.global.tuple(tensors.toArray.map(_.asInstanceOf[Tensor[?, ?]].jaxValue).toPythonProxy) val indicesAsTuple = py.Dynamic.global.tuple(ev.indices.toPythonProxy) + val jaxResult = Jax.jax_helper.zipvmap( fpy, indicesAsTuple )(jaxInputs) - Tensor(jaxResult) - + outMapper.fromPy(jaxResult) export ZipVmap.zipvmap extension [T <: Tuple: Labels, V](t: Tensor[T, V]) From 4df7732297e8236021672dcec4b75d556122ff5e Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Mon, 10 Aug 2026 12:29:40 +0200 Subject: [PATCH 2/3] Express output type more clearly. Add test case --- .../tensor/tensorops/FunctionalOps.scala | 32 +++++++++++-------- .../tensor/TensorOpsFunctionalSuite.scala | 9 ++++++ 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala index 352047d..e510267 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala @@ -17,40 +17,46 @@ import me.shadaj.scalapy.py.SeqConverters import me.shadaj.scalapy.readwrite.Reader import me.shadaj.scalapy.readwrite.Writer -trait ZipVmapResult[L: Label, FOut, MOut]: +type ZipVmapOut[L, FOut] = FOut match + case Tensor[shape, v] => Tensor[L *: shape, v] + case EmptyTuple => EmptyTuple + case h *: t => ZipVmapOut[L, h] *: ZipVmapOut[L, t] + +trait ZipVmapResult[L: Label, FOut]: type MappedOut def toPy(out: FOut): py.Dynamic - def fromPy(pyOut: py.Dynamic): MOut + def fromPy(pyOut: py.Dynamic): ZipVmapOut[L, FOut] object ZipVmapResult: // Single Tensor case - given singleTensor[L: Label, Shape <: Tuple: Labels, V]: ZipVmapResult[L, Tensor[Shape, V], Tensor[L *: Shape, V]] with + given singleTensor[L: Label, Shape <: Tuple: Labels, V]: ZipVmapResult[L, Tensor[Shape, V]] with def toPy(out: Tensor[Shape, V]): py.Dynamic = out.jaxValue def fromPy(pyOut: py.Dynamic): Tensor[L *: Shape, V] = Tensor(pyOut) // Empty Tuple case - given emptyTuple[L: Label]: ZipVmapResult[L, EmptyTuple, EmptyTuple] with + given emptyTuple[L: Label]: ZipVmapResult[L, EmptyTuple] with def toPy(out: EmptyTuple): py.Dynamic = py.Dynamic.global.tuple(Seq.empty[py.Dynamic].toPythonProxy) def fromPy(pyOut: py.Dynamic): EmptyTuple = EmptyTuple // Inductive Tuple case (Pairs, Triples, N-tuples) - given consTuple[L: Label, H, HOut, T <: Tuple, TOut <: Tuple](using - hRes: ZipVmapResult[L, H, HOut], - tRes: ZipVmapResult[L, T, TOut] - ): ZipVmapResult[L, H *: T, HOut *: TOut] with + given consTuple[L: Label, H, T <: Tuple](using + hRes: ZipVmapResult[L, H], + tRes: ZipVmapResult[L, T] + ): ZipVmapResult[L, H *: T] with def toPy(out: H *: T): py.Dynamic = val headPy = hRes.toPy(out.head) val tailSeq = tRes.toPy(out.tail).as[Seq[py.Dynamic]] py.Dynamic.global.tuple((headPy +: tailSeq).toPythonProxy) - def fromPy(pyOut: py.Dynamic): HOut *: TOut = + def fromPy(pyOut: py.Dynamic): ZipVmapOut[L, H *: T] = val seq = pyOut.as[Seq[py.Dynamic]] val h = hRes.fromPy(seq.head) val tailPy = py.Dynamic.global.tuple(seq.tail.toPythonProxy) val t = tRes.fromPy(tailPy) - h *: t + + (h *: t.asInstanceOf[Tuple]).asInstanceOf[ZipVmapOut[L, H *: T]] object FunctionalOps: @@ -86,7 +92,7 @@ object FunctionalOps: * ... * } */ - def zipvmap[L: Label, Inputs <: Tuple, FOut, MOut]( + def zipvmap[L: Label, Inputs <: Tuple, FOut]( axis: Axis[L] )( tensors: Inputs @@ -95,8 +101,8 @@ object FunctionalOps: )( f: TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]] => FOut )(using - outMapper: ZipVmapResult[L, FOut, MOut] - ): MOut = + outMapper: ZipVmapResult[L, FOut] + ): ZipVmapOut[L, FOut] = val fpy = (args: py.Dynamic) => OnError.traceStack: val tensorList = args.as[Seq[py.Dynamic]].zip(ev.shapesLabels).map: (jaxArr, labels) => diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala index 73e370f..4cff72d 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala @@ -59,6 +59,15 @@ class TensorOpsFunctionalSuite extends DimwitTest: // Each row of ta sums to 3.0, each row of tc sums to 8.0 => 11.0 per row res.shouldEqual(Tensor1(Axis[A]).fromArray(Array(11.0f, 11.0f))) + it("zipvmap2 return tuple"): + val t1 = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(0f) + val t2 = Tensor(Shape(Axis[A] -> 2, Axis[B] -> 3)).fill(1f) + val (y1, y2) = zipvmap(Axis[A])(t1, t2): + case (x1, x2) => + (x1 + x2, x1 - x2) + y1 shouldEqual (t1 + t2) + y2 shouldEqual (t1 - t2) + describe("vapply (Axis-wise application)"): def l2[L: Label](v1: Tensor1[L, Float32], v2: Tensor1[L, Float32]): Tensor0[Float32] = (v1 - v2).pow(2.0f).sum.sqrt From bcaa3fa5c88c29b3d849692ba853f264b4d85b09 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Tue, 11 Aug 2026 16:39:39 +0200 Subject: [PATCH 3/3] Add support to vmap --- .../tensor/tensorops/FunctionalOps.scala | 82 +++++++------------ .../tensor/TensorOpsFunctionalSuite.scala | 7 ++ 2 files changed, 36 insertions(+), 53 deletions(-) diff --git a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala index e510267..d882011 100644 --- a/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala +++ b/core/src/main/scala/dimwit/tensor/tensorops/FunctionalOps.scala @@ -16,49 +16,20 @@ 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 -type ZipVmapOut[L, FOut] = FOut match - case Tensor[shape, v] => Tensor[L *: shape, v] - case EmptyTuple => EmptyTuple - case h *: t => ZipVmapOut[L, h] *: ZipVmapOut[L, t] - -trait ZipVmapResult[L: Label, FOut]: - type MappedOut - def toPy(out: FOut): py.Dynamic - def fromPy(pyOut: py.Dynamic): ZipVmapOut[L, FOut] - -object ZipVmapResult: - - // Single Tensor case - given singleTensor[L: Label, Shape <: Tuple: Labels, V]: ZipVmapResult[L, Tensor[Shape, V]] with - def toPy(out: Tensor[Shape, V]): py.Dynamic = out.jaxValue - def fromPy(pyOut: py.Dynamic): Tensor[L *: Shape, V] = Tensor(pyOut) - - // Empty Tuple case - given emptyTuple[L: Label]: ZipVmapResult[L, EmptyTuple] with - def toPy(out: EmptyTuple): py.Dynamic = py.Dynamic.global.tuple(Seq.empty[py.Dynamic].toPythonProxy) - def fromPy(pyOut: py.Dynamic): EmptyTuple = EmptyTuple - - // Inductive Tuple case (Pairs, Triples, N-tuples) - given consTuple[L: Label, H, T <: Tuple](using - hRes: ZipVmapResult[L, H], - tRes: ZipVmapResult[L, T] - ): ZipVmapResult[L, H *: T] with - - def toPy(out: H *: T): py.Dynamic = - val headPy = hRes.toPy(out.head) - val tailSeq = tRes.toPy(out.tail).as[Seq[py.Dynamic]] - py.Dynamic.global.tuple((headPy +: tailSeq).toPythonProxy) - - def fromPy(pyOut: py.Dynamic): ZipVmapOut[L, H *: T] = - val seq = pyOut.as[Seq[py.Dynamic]] - val h = hRes.fromPy(seq.head) - val tailPy = py.Dynamic.global.tuple(seq.tail.toPythonProxy) - val t = tRes.fromPy(tailPy) +object FunctionalOps: - (h *: t.asInstanceOf[Tuple]).asInstanceOf[ZipVmapOut[L, H *: T]] + type PrependAxes[Axes <: Tuple, FOut] = Axes match + case EmptyTuple => FOut + case h *: t => PrependAxis[h, PrependAxes[t, FOut]] -object FunctionalOps: + 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] object ZipVmap: @@ -101,8 +72,9 @@ object FunctionalOps: )( f: TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]] => FOut )(using - outMapper: ZipVmapResult[L, FOut] - ): ZipVmapOut[L, FOut] = + toPyTree: TensorTree[FOut], + fromPyTree: TensorTree[PrependAxis[L, FOut]] + ): PrependAxis[L, FOut] = val fpy = (args: py.Dynamic) => OnError.traceStack: val tensorList = args.as[Seq[py.Dynamic]].zip(ev.shapesLabels).map: (jaxArr, labels) => @@ -110,7 +82,7 @@ object FunctionalOps: val inputTuple = Tuple.fromArray(tensorList.toArray) val result = f(inputTuple.asInstanceOf[TensorsOf[ev.RemainingAxes, ValuesOf[Inputs]]]) - outMapper.toPy(result) + toPyTree.toPyTree(result) val jaxInputs = py.Dynamic.global.tuple(tensors.toArray.map(_.asInstanceOf[Tensor[?, ?]].jaxValue).toPythonProxy) val indicesAsTuple = py.Dynamic.global.tuple(ev.indices.toPythonProxy) @@ -120,7 +92,7 @@ object FunctionalOps: indicesAsTuple )(jaxInputs) - outMapper.fromPy(jaxResult) + fromPyTree.fromPyTree(jaxResult) export ZipVmap.zipvmap extension [T <: Tuple: Labels, V](t: Tensor[T, V]) @@ -133,13 +105,16 @@ object FunctionalOps: * @param f A function that takes a tuple of tensors (with the specified axis removed) and returns a new tensor. * @return A new tensor resulting from applying `f` to the zipped tensors. */ - def zipvmap[L: Label, T2 <: Tuple, OutShape <: Tuple: Labels, OutV](axis: Axis[L])( + def zipvmap[L: Label, T2 <: Tuple, FOut](axis: Axis[L])( other: Tensor[T2, V] )(using ev: SharedAxisRemover[(T, T2), L] )( - f: TensorsOf[ev.RemainingAxes, (V, V)] => Tensor[OutShape, OutV] - ): Tensor[L *: OutShape, OutV] = + f: TensorsOf[ev.RemainingAxes, (V, V)] => FOut + )(using + toPyTree: TensorTree[FOut], + fromPyTree: TensorTree[PrependAxis[L, FOut]] + ): PrependAxis[L, FOut] = ZipVmap.zipvmap(axis)(t, other)(f) /** Vectorized mapping over a specified axis of the tensor. @@ -148,22 +123,23 @@ object FunctionalOps: * @param f A function that takes a tensor with the specified axis removed and returns a new tensor. * @return A new tensor resulting from applying `f` to each slice along the specified axis. */ - def vmap[VmapAxis: Label, OuterShape <: Tuple: Labels, V2]( + def vmap[VmapAxis: Label, FOut]( axis: Axis[VmapAxis] )(using ev: AxisRemover[T, VmapAxis] )( - f: Tensor[ev.RemainingAxes, V] => Tensor[OuterShape, V2] + f: Tensor[ev.RemainingAxes, V] => FOut )(using + toPyTree: TensorTree[FOut], + fromPyTree: TensorTree[PrependAxis[VmapAxis, FOut]], labels: Labels[ev.RemainingAxes] - ): Tensor[VmapAxis *: OuterShape, V2] = + ): PrependAxis[VmapAxis, FOut] = val fpy = (jxpr: Jax.PyDynamic) => OnError.traceStack: val innerTensor = Tensor[ev.RemainingAxes, V](jxpr) val result = f(innerTensor) - result.jaxValue - - Tensor(Jax.jax_helper.vmap(fpy, ev.index)(t.jaxValue)) + toPyTree.toPyTree(result) + fromPyTree.fromPyTree(Jax.jax_helper.vmap(fpy, ev.index)(t.jaxValue)) /** Apply a function independently to each 1D slice along a labeled axis. * diff --git a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala index 4cff72d..0278b1f 100644 --- a/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala +++ b/core/src/test/scala/dimwit/tensor/TensorOpsFunctionalSuite.scala @@ -22,6 +22,13 @@ class TensorOpsFunctionalSuite extends DimwitTest: val res = t2.vmap(Axis[A])(_.sum) res shouldEqual Tensor1(Axis[A]).fromArray(Array(3.0f, 7.0f)) + 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) + y1 shouldEqual (t +! 5f) + y2 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))