To numpy array

[benikm91]

toPyTree and fromPyTree are good for casting a TensorTree to a Python tree for execution of operations.

However, when serializing a tensor tree to disc (e.g., TensorTreeFormat.Pickle), the function fails for a TensorTree with a Random.Key. Instead of a custom logic in TensorTreeFormat.Pickle I propose adding a toNumpyTree and fromNumPyTree for serialization.

case class TrainState(params: UNet.Params, lastCost: Tensor0[Float32], key: Random.Key)

Key needs custom logic to be cast from DimWit/JAX to Numpy. We may could put that into TensorTreeFormat.Pickle toHost, however this is a more general solution if we need it elsewhere or if other special cases appear.

given TensorTree[Key] with
     ...
      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])))

FYI: I tried storing a checkpoint and restarting from it. With same data samples the exact same train trajectory was followed, so the key write/read works successfully :)

[marcelluethi]

That looks good to me.

[benikm91]

Should be merged after #141