tree_to_jax

t3toolbox.backend.common.tree_to_jax(T)

def tree_to_jax(T):

Move every array leaf of a pytree (nested tuples/lists of arrays) onto jax, preserving the tree structure; leaves already jax are a no-op (jnp.asarray). Numpy float64 leaves become jax float32 unless jax x64 is enabled – the caller opts into jax precision. Requires jax – guard on jax_available at the call site (jnp is only bound when jax is importable).