OpTree: Optimized PyTree Utilities
Visualize, create, and operate on pytrees in the most intuitive way possible.
A small library for creating and manipulating custom JAX Pytree classes
A stateful pytree library for training neural networks.
io+hdf5+lazy-loading+processing-pipelines for EHRs and generic PyTrees for JAX friendly tasks, mainly ML