To numpy array

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 {}
 

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

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

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

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)