PyTree Utilities

February 9, 2026 · View on GitHub

JAX pytree manipulation utilities.

See Also

Utils Overview - All utilities
NNX Best Practices - JAX patterns
Checkpointing - Pytree serialization
Sharding - Pytree sharding

::: datarax.utils.pytree_utils