Introduce TensorTreeIO and reorganize tensor tree

AGENTS.md

@@ -758,7 +758,8 @@ Use **case classes** to group parameters. DimWit automatically derives `TensorTr
 
 ```scala
 import dimwit.*
-import dimwit.autodiff.{TensorTree, FloatTree, Autodiff}
+import dimwit.autodiff.{Autodiff}
+import dimwit.tensortree.{TensorTree, FloatTree}
 
 trait Feature derives Label
 trait Hidden derives Label
@@ -1283,23 +1284,23 @@ val wrong = Autodiff.grad(nonScalar)  // Use jacobian instead
 //  [Input, V]
 //   (f: Input => dimwit.tensor.Tensor0[V])
 //     (using evidence$1: dimwit.tensor.TensorOps.IsFloating[V],
-//       inTree: dimwit.autodiff.TensorTree[Input], outTree:
-//       dimwit.autodiff.TensorTree[dimwit.tensor.Tensor0[V]]): Input =>
+//       inTree: dimwit.tensortree.TensorTree[Input], outTree:
+//       dimwit.tensortree.TensorTree[dimwit.tensor.Tensor0[V]]): Input =>
 //       dimwit.autodiff.Grad[Input]
 //  [T1, T2, T3, V²]
 //   (f: (T1, T2, T3) => dimwit.tensor.Tensor0[V²])
 //     (using evidence$1²: dimwit.tensor.TensorOps.IsFloating[V²],
-//       t1Tree: dimwit.autodiff.TensorTree[T1],
-//       t2Tree: dimwit.autodiff.TensorTree[T2],
-//       t3Tree: dimwit.autodiff.TensorTree[T3], outTree²:
-//       dimwit.autodiff.TensorTree[dimwit.tensor.Tensor0[V²]]): (T1, T2, T3) =>
+//       t1Tree: dimwit.tensortree.TensorTree[T1],
+//       t2Tree: dimwit.tensortree.TensorTree[T2],
+//       t3Tree: dimwit.tensortree.TensorTree[T3], outTree²:
+//       dimwit.tensortree.TensorTree[dimwit.tensor.Tensor0[V²]]): (T1, T2, T3) =>
 //       dimwit.autodiff.Grad[(T1, T2, T3)]
 //  [T1², T2², V³]
 //   (f: (T1², T2²) => dimwit.tensor.Tensor0[V³])
 //     (using evidence$1³: dimwit.tensor.TensorOps.IsFloating[V³],
-//       t1Tree²: dimwit.autodiff.TensorTree[T1²],
-//       t2Tree²: dimwit.autodiff.TensorTree[T2²], outTree³:
-//       dimwit.autodiff.TensorTree[dimwit.tensor.Tensor0[V³]]): (T1², T2²) =>
+//       t1Tree²: dimwit.tensortree.TensorTree[T1²],
+//       t2Tree²: dimwit.tensortree.TensorTree[T2²], outTree³:
+//       dimwit.tensortree.TensorTree[dimwit.tensor.Tensor0[V³]]): (T1², T2²) =>
 //       dimwit.autodiff.Grad[(T1², T2²)]
 // match arguments (dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.Float32] =>
 //   dimwit.tensor.Tensor1[MdocApp12.this.A, dimwit.Float32])

core/src/main/scala/dimwit/MemoryHelper.scala

@@ -1,6 +1,6 @@
 package dimwit
 
-import dimwit.autodiff.TensorTree
+import dimwit.tensortree.TensorTree
 import me.shadaj.scalapy.py
 
 private[dimwit] object MemoryHelper:

core/src/main/scala/dimwit/autodiff/Autodiff.scala

@@ -6,6 +6,7 @@ import dimwit.tensor.Tensor
 import dimwit.tensor.Tensor0
 import dimwit.tensor.TensorOps.IsFloating
 import dimwit.tensor.TupleHelpers.PrimeConcatType
+import dimwit.tensortree.TensorTree
 import me.shadaj.scalapy.py
 
 object Autodiff:

core/src/main/scala/dimwit/jax/Jit.scala

@@ -1,7 +1,7 @@
 package dimwit.jax
 
 import dimwit.OnError
-import dimwit.autodiff.TensorTree
+import dimwit.tensortree.TensorTree
 import dimwit.jax.Jax
 import dimwit.jax.Jax.PyDynamic
 import me.shadaj.scalapy.py

core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala

@@ -1,9 +1,11 @@
 package dimwit.optimizer
 
 import dimwit.*
+
 import dimwit.Conversions.given
-import dimwit.autodiff.FloatTree.*
-import dimwit.autodiff.FloatTree.ops.*
+import dimwit.tensortree.*
+import dimwit.tensortree.FloatTree.*
+import dimwit.tensortree.FloatTree.ops.*
 import dimwit.autodiff.*
 
 /** Gradient optimizer interface with functional state management.

core/src/main/scala/dimwit/package.scala

@@ -84,7 +84,9 @@ package object dimwit:
   // Export devices
   export dimwit.hardware.Device
   // Export automatic differentiation
-  export dimwit.autodiff.{Autodiff, TensorTree, FloatTree, Grad}
+  export dimwit.autodiff.{Autodiff, Grad}
+  // Export tensor trees
+  export dimwit.tensortree.{TensorTree, TensorTreeIO, TensorTreeFormat, FloatTree}
   // Export Just-in-Time compilation
   export dimwit.jax.Jit.{jit, jitDonating, jitDonatingUnsafe}
   export dimwit.jax.EagerCleanup.eagerCleanup

core/src/main/scala/dimwit/python/PyBridge.scala

@@ -1,7 +1,7 @@
 package dimwit.python
 
 import dimwit.OnError
-import dimwit.autodiff.TensorTree
+import dimwit.tensortree.TensorTree
 import dimwit.jax.Jax
 import dimwit.tensor.*
 import me.shadaj.scalapy.py

core/src/main/scala/dimwit/random/Random.scala

@@ -1,6 +1,6 @@
 package dimwit.random
 
-import dimwit.autodiff.TensorTree
+import dimwit.tensortree.TensorTree
 import dimwit.jax.Jax
 import dimwit.python.PyBridge.liftPyTensor
 import dimwit.tensor.DType.Int32

core/src/main/scala/dimwit/autodiff/FloatTree.scala → core/src/main/scala/dimwit/tensortree/FloatTree.scala

@@ -1,4 +1,4 @@
-package dimwit.autodiff
+package dimwit.tensortree
 
 import dimwit.tensor.TensorOps.*
 import dimwit.tensor.*

core/src/main/scala/dimwit/autodiff/TensorTree.scala → core/src/main/scala/dimwit/tensortree/TensorTree.scala

@@ -1,4 +1,4 @@
-package dimwit.autodiff
+package dimwit.tensortree
 
 import dimwit.jax.Jax
 import dimwit.tensor.*

core/src/main/scala/dimwit/tensortree/TensorTreeFormat.scala

@@ -0,0 +1,40 @@
+package dimwit.tensortree
+
+import dimwit.jax.Jax
+import me.shadaj.scalapy.py
+
+import java.nio.file.Path
+
+/** Provides an interface for reading and writing tensor trees to and from disk in various formats.
+  */
+trait TensorTreeFormat:
+  def write[P](p: P, path: Path)(using tt: TensorTree[P]): Unit
+  def read[P](path: Path)(using tt: TensorTree[P]): P
+
+object TensorTreeFormat:
+
+  /** Pickle format for saving and loading tensor trees.
+    *
+    * Arrays are moved off-device to numpy before pickling, and re-materialized
+    * as JAX arrays on load, so files remain portable across host/GPU/TPU.
+    * `jax.tree_util.tree_map` handles the tuple/list/None nesting natively, so
+    * no recursive walk is needed here.
+    */
+  object Pickle extends TensorTreeFormat:
+    private lazy val pickle = py.module("pickle")
+    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 file = builtins.open(path.toAbsolutePath().toString(), "wb").as[py.Dynamic]
+      try pickle.dump(numpyTree, file)
+      finally file.close()
+
+    def read[P](path: Path)(using tt: TensorTree[P]): P =
+      val file = builtins.open(path.toAbsolutePath().toString(), "rb").as[py.Dynamic]
+      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))

core/src/main/scala/dimwit/tensortree/TensorTreeIO.scala

@@ -0,0 +1,27 @@
+package dimwit.tensortree
+
+import java.nio.file.Path
+
+/** Provides methods to save and load tensor trees to and from disk.
+  * The default format is pickle, but other formats can be specified if needed.
+  */
+object TensorTreeIO:
+
+  /** Saves a tensor tree to disk in the specified format.
+    *
+    * @param p The tensor tree structure (e.g., model parameters) to be saved.
+    * @param path The file path where the tensor tree will be saved.
+    * @param format The format in which to save the tensor tree (default is pickle).
+    */
+  def save[P](p: P, path: Path, format: TensorTreeFormat = TensorTreeFormat.Pickle)(using tt: TensorTree[P]): Unit =
+    format.write(p, path)
+
+  /** Loads a tensor tree from disk in the specified format.
+    *
+    * @tparam P The type of the tensor tree structure to be loaded.
+    * @param path The file path from which to load the tensor tree.
+    * @param format The format in which to load the tensor tree (default is pickle).
+    * @return The loaded tensor tree structure.
+    */
+  def load[P](path: Path, format: TensorTreeFormat = TensorTreeFormat.Pickle)(using tt: TensorTree[P]): P =
+    format.read(path)

core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala

@@ -2,8 +2,8 @@ package dimwit.optimizer
 
 import dimwit.*
 import dimwit.Conversions.given
-import dimwit.autodiff.FloatTree.*
-import dimwit.autodiff.FloatTree.ops.*
+import dimwit.tensortree.FloatTree.*
+import dimwit.tensortree.FloatTree.ops.*
 import dimwit.autodiff.*
 
 class GradientOptimizerSuite extends DimwitTest:

core/src/test/scala/dimwit/python/PyWrapSuite.scala

@@ -77,7 +77,7 @@ class PyWrapSuite extends DimwitTest:
       // Call through a Python lambda that just forwards to pyFn
       val caller = py.eval("lambda fn, x: fn(x)")
       val input = Tensor1(Axis[A]).fromArray(Array(2f, 3f, 4f))
-      val pyResult = caller(pyFn, dimwit.autodiff.TensorTree[Tensor1[A, Float32]].toPyTree(input))
-      val result = dimwit.autodiff.TensorTree[Tensor1[A, Float32]].fromPyTree(pyResult)
+      val pyResult = caller(pyFn, dimwit.tensortree.TensorTree[Tensor1[A, Float32]].toPyTree(input))
+      val result = dimwit.tensortree.TensorTree[Tensor1[A, Float32]].fromPyTree(pyResult)
 
       result should approxEqual(Tensor1(Axis[A]).fromArray(Array(6f, 9f, 12f)))

core/src/test/scala/dimwit/autodiff/FloatTensorTreeSuite.scala → core/src/test/scala/dimwit/tensortree/FloatTensorTreeSuite.scala

@@ -1,9 +1,9 @@
-package dimwit.autodiff
+package dimwit.tensortree
 
 import dimwit.*
 import dimwit.Conversions.given
-import dimwit.autodiff.FloatTree.*
-import dimwit.autodiff.FloatTree.ops.*
+import dimwit.tensortree.FloatTree.*
+import dimwit.tensortree.FloatTree.ops.*
 
 class FloatTensorTreeSuite extends DimwitTest:
 

core/src/test/scala/dimwit/autodiff/PyTreeSuite.scala → core/src/test/scala/dimwit/tensortree/PyTreeSuite.scala

@@ -1,4 +1,4 @@
-package dimwit.autodiff
+package dimwit.tensortree
 
 import dimwit.*
 import dimwit.jax.Jax

core/src/test/scala/dimwit/tensortree/TensorTreeIOSuite.scala

@@ -0,0 +1,44 @@
+package dimwit.tensortree
+
+import dimwit.*
+
+import java.nio.file.Files
+
+class TensorTreeIOSuite extends DimwitTest:
+
+  describe("save and load"):
+    it("round-trips a 1-level case class"):
+      case class Params(w1: Tensor1[A, Float32], b1: Tensor0[Int32])
+      val params = Params(
+        Tensor1(Axis[A]).fromArray(Array(0.1f, 0.2f, 0.3f)),
+        Tensor0(42)
+      )
+      val path = Files.createTempFile("tensortree-io", ".pkl")
+      try
+        TensorTreeIO.save(params, path)
+        val restored = TensorTreeIO.load[Params](path)
+        restored.w1 should approxEqual(params.w1)
+        restored.b1 should equal(params.b1)
+      finally Files.deleteIfExists(path)
+
+    it("round-trips nested structures (lists and tuples)"):
+      case class Model(layers: List[Tensor0[Float32]], extra: (Tensor0[Int32], Tensor0[Int32]))
+      val params = Model(
+        List(Tensor0(1.0f), Tensor0(2.0f), Tensor0(3.0f)),
+        (Tensor0(3), Tensor0(4))
+      )
+      val path = Files.createTempFile("tensortree-io-nested", ".pkl")
+      try
+        TensorTreeIO.save(params, path)
+        val restored = TensorTreeIO.load[Model](path)
+        restored.layers should equal(params.layers)
+        restored.extra should equal(params.extra)
+      finally Files.deleteIfExists(path)
+
+    it("round-trips an empty (Unit) tree"):
+      val path = Files.createTempFile("tensortree-io-unit", ".pkl")
+      try
+        TensorTreeIO.save[Unit]((), path)
+        val restored = TensorTreeIO.load[Unit](path)
+        restored should equal(())
+      finally Files.deleteIfExists(path)

core/src/test/scala/dimwit/autodiff/TensorTreeSuite.scala → core/src/test/scala/dimwit/tensortree/TensorTreeSuite.scala

@@ -1,4 +1,4 @@
-package dimwit.autodiff
+package dimwit.tensortree
 
 import dimwit.*
 

examples/src/main/scala/complex/VariationalAutoencoder.scala

@@ -2,7 +2,7 @@ package examples.complex.vae
 
 import dimwit.Conversions.given
 import dimwit.*
-import dimwit.autodiff.FloatTree.*
+import dimwit.tensortree.FloatTree.*
 import dimwit.autodiff.*
 import dimwit.nn.ActivationFunctions.relu
 import dimwit.nn.ActivationFunctions.sigmoid

mdocs/AGENTS.md

@@ -573,7 +573,8 @@ Use **case classes** to group parameters. DimWit automatically derives `TensorTr
 
 ```scala mdoc:reset:silent
 import dimwit.*
-import dimwit.autodiff.{TensorTree, FloatTree, Autodiff}
+import dimwit.autodiff.{Autodiff}
+import dimwit.tensortree.{TensorTree, FloatTree}
 
 trait Feature derives Label
 trait Hidden derives Label