diff --git a/core/src/main/scala/dimwit/autodiff/Grad.scala b/core/src/main/scala/dimwit/autodiff/Grad.scala index f4d862c..8e67f5e 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)) + // TreeOf witness for gradient math (++, --, scale, etc.) // given [T, V: IsFloating](using TreeOf[T, V]): TreeOf[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)