Change zipvmap to allow returning tuples of tensors

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,33 +86,35 @@ 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) =>
             Tensor(jaxArr)(using LabelsImpl(labels))
 
           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])