orbax

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

Training, testing and CI infrastructure for JAX/Flax NNX: JAX runtime setup, device meshes and SPMD sharding, optax optimizers, an Orbax checkpoint store, early stopping, experiment logging, GPU compute jobs, fresh-interpreter test tools, and reusable CI gates and actions.

Reproducible JAX/Keras 3 + Orbax runtime and TPU v5e-8 production qualification for Microsoft Mage-Flow-Turbo, with bilingual docs, public model artifacts, and acceptance evidence.

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.

Production-qualified JAX/Keras 3 port of Microsoft Mage-Flow-Edit-Turbo for Kaggle TPU v5e-8, with standalone JAX/Orbax artifacts, public acceptance evidence, deterministic validation at 512/768/1024, reproducibility assets, licensing/provenance docs, and a checksummed v1.0.0 release bundle.

flax.nnx implementation of the UvA Deep Learning Tutorials