Substrax

The training plumbing a JAX/Flax NNX project otherwise writes for itself, written once: configure the JAX process, find the devices and lay arrays out across them, save and restore state, decide when to stop, log the run, and send it to a GPU.

Every Avitai JAX package imports this layer instead of carrying its own copy, so each concern has one implementation and one test suite. Substrax imports none of those packages, so it works on its own.

The dependency chain: substrax → calibrax → datarax → artifex → opifex.

Subpackage	What it owns
substrax.runtime	JaxRuntime process settings rendered as the environment of a process that has not imported jax, or applied to the current one; XLA flags merged by name; test-run device emulation; entry-point logging
substrax.artifacts	Output directories resolved from an argument, AVITAI_OUTPUT_DIR or a per-run temporary directory, never the working tree
substrax.rng	Keys from an explicit owner (key_from, no default seed), streams derived from a seed by name (rngs_from_seed), split_key and the interpreter-stable fold_in_name
substrax.optim	OptimizerConfig in optax's terms, create_transformation and create_optimizer over optax with the schedule as the base learning rate and the weight-decay filter as a static mask, current_learning_rate read on device, switch_at from one transformation to another at a step, update_with_line_search for L-BFGS-style optimizers, with_strong_state_types
substrax.testing	Opt-in test infrastructure: fresh-interpreter runs with a chosen JAX configuration and typed JSON reports (last_json_as), trace, compile and work counters (TraceCounter, substrax.testing.compiles, counted_calls), value-asserting NNX gradient checks (substrax.testing.gradients), example runs, and a pytest plugin with x64, devices and accelerator markers and jax configuration isolation
substrax.devices	detect_devices() (platform, device kind, count), DeviceLike and visible_devices() for typed device parameters, device placement, the batch-size recommendation table
substrax.mesh	Device meshes with Auto axes by default, mesh rules and partition-spec helpers, sharding strategies (data, FSDP, tensor, pipeline, multi-dimensional) on flax.nnx.spmd
substrax.spmd	Data-parallel sharding, batch placement from a sharding or a prefix of the batch, spmd_train_step, and metric collectives
substrax.checkpoint	One CheckpointStore protocol and one Orbax implementation, OrbaxCheckpointStore, over CheckpointManager
substrax.callbacks	The training-callback protocol, CallbackList, BestMetricTracker, EarlyStopping and EarlyStoppingCallback
substrax.tracking	Step-wise experiment logging with console, file, Weights & Biases and MLflow backends
substrax.compute	A project's jobs run on a compute backend: the job spec, the worker every backend runs, the ComputeBackend protocol with local, modal and skypilot backends found through entry points, the accelerator, CPU and memory a job requests, a manifest recording the devices a run got, and the substrax-compute command
substrax.typing, substrax.nnx_typing, substrax.records, substrax.examples	The shared type aliases (PyTree, JsonValue, CheckpointState; NnxState and PathEntry over jax and flax), typed reading and writing of JSON records, and the listing of a repository's example scripts

Each subpackage has an API page.

Installation

uv add substrax

Substrax requires Python 3.12 or later (CI runs 3.12 and 3.13), jax and jaxlib 0.11.1, flax>=0.12.9, optax>=0.2.8, orbax-checkpoint>=0.11.33, etils[epath]>=1.14.0, numpy>=1.24,<2.6 and pydantic>=2.10. jax is capped below 0.11.2 until a flax release works with it (pyproject.toml records why). The cuda12 and metal extras select the JAX backend.