From 4c26e711b354c513b751add806583fe06a0b3a24 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Fri, 7 Aug 2026 13:58:09 +0200 Subject: [PATCH 1/2] 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 2447962a99a3375622904e3593a471319c2a9658 Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Fri, 7 Aug 2026 20:04:45 +0200 Subject: [PATCH 2/2] Add toNumpyTree fromNumpyTree for serialization to numpy --- .../src/main/scala/dimwit/autodiff/Grad.scala | 4 ++ .../src/main/scala/dimwit/random/Random.scala | 8 ++++ .../scala/dimwit/tensortree/TensorTree.scala | 48 +++++++++++++++++++ .../dimwit/tensortree/TensorTreeFormat.scala | 6 +-- 4 files changed, 62 insertions(+), 4 deletions(-) diff --git a/core/src/main/scala/dimwit/autodiff/Grad.scala b/core/src/main/scala/dimwit/autodiff/Grad.scala index 27f9964..bf9b7ce 100644 --- a/core/src/main/scala/dimwit/autodiff/Grad.scala +++ b/core/src/main/scala/dimwit/autodiff/Grad.scala @@ -51,6 +51,10 @@ object Grad: def fromPyTree(pyVal: Jax.PyAny): Grad[T] = Grad(ev.fromPyTree(pyVal)) + def toNumpyTree(g: Grad[T]): Jax.PyAny = ev.toNumpyTree(g) + + def fromNumpyTree(pyVal: Jax.PyAny): Grad[T] = Grad(ev.fromNumpyTree(pyVal)) + // FloatTree witness for gradient math (++, --, scale, etc.) // given [T, V: IsFloating](using FloatTree[T, V]): FloatTree[Grad[T], V] with {} diff --git a/core/src/main/scala/dimwit/random/Random.scala b/core/src/main/scala/dimwit/random/Random.scala index 3e93222..bdd89bc 100644 --- a/core/src/main/scala/dimwit/random/Random.scala +++ b/core/src/main/scala/dimwit/random/Random.scala @@ -114,6 +114,14 @@ object Random: def fromPyTree(pyVal: Jax.PyAny): Key = Key(pyVal.as[Jax.PyDynamic]) + def toNumpyTree(p: Key): Jax.PyAny = + // Extract key data for numpy serialization using key_data + Jax.np.asarray(Jax.jax.device_get(Jax.jax.random.key_data(p.jaxKey))) + + def fromNumpyTree(pyVal: Jax.PyAny): Key = + // Reconstruct key from numpy array using wrap_key_data + Key(Jax.jax.random.wrap_key_data(Jax.jnp.asarray(pyVal.as[Jax.PyDynamic]))) + /** Create a random key from an integer seed */ def apply(seed: Long): Key = Key(Jax.jrandom.key(seed)) diff --git a/core/src/main/scala/dimwit/tensortree/TensorTree.scala b/core/src/main/scala/dimwit/tensortree/TensorTree.scala index 9467e45..fb11cb4 100644 --- a/core/src/main/scala/dimwit/tensortree/TensorTree.scala +++ b/core/src/main/scala/dimwit/tensortree/TensorTree.scala @@ -60,6 +60,16 @@ trait TensorTree[P]: */ def fromPyTree(py: Jax.PyAny): P + /** Convert the structure p to a tree representation of numpy arrays. + * While toPyTree is for in-memory representation, toNumpyTree is for saving to disk or sending over the network. + */ + def toNumpyTree(p: P): Jax.PyAny + + /** Convert a tree representation of numpy arrays back to the structure P. + * While fromPyTree is for in-memory representation, fromNumpyTree is for loading from disk or receiving over the network. + */ + def fromNumpyTree(pyVal: Jax.PyAny): P + object TensorTree: // extends TensorTreeLowPriority: def apply[P](using pt: TensorTree[P]): TensorTree[P] = pt @@ -121,6 +131,9 @@ object TensorTree: // extends TensorTreeLowPriority: def toPyTree(p: Tensor[Q, V]): Jax.PyAny = p.jaxValue def fromPyTree(pyVal: Jax.PyAny): Tensor[Q, V] = Tensor(pyVal.as[Jax.PyDynamic]) + def toNumpyTree(p: Tensor[Q, V]): Jax.PyAny = Jax.np.asarray(Jax.jax.device_get(p.jaxValue)) + def fromNumpyTree(pyVal: Jax.PyAny): Tensor[Q, V] = Tensor(Jax.jnp.asarray(pyVal)) + /** Tensor tree instance for an empty tree. This can be useful * for example for optimizers that don't have internal state */ @@ -133,6 +146,8 @@ object TensorTree: // extends TensorTreeLowPriority: def zipMap(p1: Unit, p2: Unit, f: [T <: Tuple, V] => (Labels[T]) ?=> ((Tensor[T, V], Tensor[T, V]) => Tensor[T, V])): Unit = () def toPyTree(p: Unit): Jax.PyAny = py.Dynamic.global.None def fromPyTree(pyVal: Jax.PyAny): Unit = () + def toNumpyTree(p: Unit): Jax.PyAny = py.Dynamic.global.None + def fromNumpyTree(pyVal: Jax.PyAny): Unit = () /** Instance for a tuple of two tensors */ given tupleInstance[P1, P2](using t1: TensorTree[P1], t2: TensorTree[P2]): TensorTree[(P1, P2)] with @@ -167,6 +182,13 @@ object TensorTree: // extends TensorTreeLowPriority: val pyTuple = pyVal.as[py.Dynamic] (t1.fromPyTree(pyTuple.bracketAccess(0)), t2.fromPyTree(pyTuple.bracketAccess(1))) + def toNumpyTree(p: (P1, P2)): Jax.PyAny = + py.Dynamic.global.tuple(Seq(t1.toNumpyTree(p._1), t2.toNumpyTree(p._2)).toPythonProxy) + + def fromNumpyTree(pyVal: Jax.PyAny): (P1, P2) = + val pyTuple = pyVal.as[py.Dynamic] + (t1.fromNumpyTree(pyTuple.bracketAccess(0)), t2.fromNumpyTree(pyTuple.bracketAccess(1))) + /** Instance for a list of tensor trees */ given listInstance[P](using tp: TensorTree[P]): TensorTree[List[P]] with @@ -201,6 +223,15 @@ object TensorTree: // extends TensorTreeLowPriority: val len = py.Dynamic.global.len(pyList).as[Int] List.tabulate(len)(i => tp.fromPyTree(pyList.bracketAccess(i))) + def toNumpyTree(l: List[P]): Jax.PyAny = + val pyItems = l.map(a => tp.toNumpyTree(a)) + py.Dynamic.global.list(pyItems.toPythonProxy) + + def fromNumpyTree(pyVal: Jax.PyAny): List[P] = + val pyList = pyVal.as[py.Dynamic] + val len = py.Dynamic.global.len(pyList).as[Int] + List.tabulate(len)(i => tp.fromNumpyTree(pyList.bracketAccess(i))) + given namedTupleInstance[N <: Tuple, V <: Tuple](using tt: TensorTree[V]): TensorTree[NamedTuple[N, V]] with def map(p: NamedTuple[N, V], f: [T <: Tuple, V2] => (Labels[T]) ?=> (Tensor[T, V2] => Tensor[T, V2])): NamedTuple[N, V] = tt.map(p.toTuple, f) @@ -226,6 +257,12 @@ object TensorTree: // extends TensorTreeLowPriority: def fromPyTree(pyVal: Jax.PyAny): NamedTuple[N, V] = tt.fromPyTree(pyVal) + def toNumpyTree(p: NamedTuple[N, V]): Jax.PyAny = + tt.toNumpyTree(p.toTuple) + + def fromNumpyTree(pyVal: Jax.PyAny): NamedTuple[N, V] = + tt.fromNumpyTree(pyVal) + /** automatically derive a TensorTree instance for any case class (or product type) * whose fields all have TensorTree instances. */ @@ -299,3 +336,14 @@ object TensorTree: // extends TensorTreeLowPriority: val elems = instances.zipWithIndex.map: (tc, index) => tc.fromPyTree(pyTuple.bracketAccess(index)) m.fromProduct(Tuple.fromArray(elems.map(_.asInstanceOf[Object]).toArray)) + + def toNumpyTree(p: P): Jax.PyAny = + val pyTreeElems = p.productIterator.toList.zip(instances).map: + case (field, tc) => tc.toNumpyTree(field) + py.Dynamic.global.tuple(pyTreeElems.toPythonProxy) + + def fromNumpyTree(pyVal: Jax.PyAny): P = + val pyTuple = pyVal.as[py.Dynamic] + val elems = instances.zipWithIndex.map: (tc, index) => + tc.fromNumpyTree(pyTuple.bracketAccess(index)) + m.fromProduct(Tuple.fromArray(elems.map(_.asInstanceOf[Object]).toArray)) diff --git a/core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala b/core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala index 62ceee9..028fbf7 100644 --- a/core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala +++ b/core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala @@ -25,8 +25,7 @@ object TensorTreeFormat: private lazy val builtins = py.module("builtins") def write[P](p: P, path: Path)(using tt: TensorTree[P]): Unit = - val toHost = (x: Jax.PyDynamic) => Jax.np.asarray(Jax.jax.device_get(x)) - val numpyTree = Jax.jax.tree_util.tree_map(toHost, tt.toPyTree(p)) + val numpyTree = tt.toNumpyTree(p) val file = builtins.open(path.toAbsolutePath().toString(), "wb").as[py.Dynamic] try pickle.dump(numpyTree, file) finally file.close() @@ -36,5 +35,4 @@ object TensorTreeFormat: val numpyTree = try pickle.load(file).as[py.Dynamic] finally file.close() - val toDevice = (x: Jax.PyDynamic) => Jax.jnp.asarray(x) - tt.fromPyTree(Jax.jax.tree_util.tree_map(toDevice, numpyTree)) + tt.fromNumpyTree(numpyTree)