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])