orbax

Differentiable data pipelines for JAX/Flax NNX: HuggingFace, TFDS and ArrayRecord sources, augmentation stages, scan-based epochs, exact mid-epoch resume.

Tungsten Alpha: A JAX-native, 7-layer Adelic-Riemannian Operating System. Implements resurgent gradient flow through discrete logic and Levi-Civita parallel transport for decentralized XLA clusters. Hardened for zero-jitter, systolic array execution.

flax.nnx implementation of the UvA Deep Learning Tutorials

JAX/Flax NNX training infrastructure: device meshes and SPMD sharding, an Orbax checkpoint store, early stopping and callbacks, W&B and MLflow logging.