Metadata-Version: 2.5
Name: substrax
Version: 0.1.9
Summary: JAX/Flax NNX training infrastructure: device meshes and SPMD sharding, an Orbax checkpoint store, early stopping and callbacks, W&B and MLflow logging
Project-URL: Changelog, https://github.com/avitai/substrax/blob/main/CHANGELOG.md
Project-URL: Documentation, https://substrax.readthedocs.io
Project-URL: Repository, https://github.com/avitai/substrax
Author-email: Mahdi Shafiei <mahdi@avitai.bio>
License: MIT License
        
        Copyright (c) 2026 Mahdi Shafiei
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Keywords: checkpointing,distributed,flax,infrastructure,jax,sharding
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Natural Language :: English
Classifier: Operating System :: MacOS
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.12
Requires-Dist: flax>=0.12.9
Requires-Dist: jax>=0.11.1
Requires-Dist: jaxlib>=0.11.1
Requires-Dist: numpy<2.6.0,>=1.24
Requires-Dist: optax>=0.2.8
Requires-Dist: orbax-checkpoint>=0.11.33
Provides-Extra: all
Requires-Dist: bandit[toml]>=1.8.6; extra == 'all'
Requires-Dist: build>=1.0.3; extra == 'all'
Requires-Dist: coverage>=7; extra == 'all'
Requires-Dist: deptry>=0.23; extra == 'all'
Requires-Dist: griffe>=1.7.3; extra == 'all'
Requires-Dist: hypothesis>=6.100; extra == 'all'
Requires-Dist: import-linter>=2.0; extra == 'all'
Requires-Dist: interrogate>=1.5; extra == 'all'
Requires-Dist: jax[cuda12]>=0.11.1; extra == 'all'
Requires-Dist: matplotlib>=3.8; extra == 'all'
Requires-Dist: mkdocs-material>=9.7.7; extra == 'all'
Requires-Dist: mkdocs<2.0,>=1.6.1; extra == 'all'
Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'all'
Requires-Dist: mkdocstrings>=0.28.3; extra == 'all'
Requires-Dist: mlflow>=2; extra == 'all'
Requires-Dist: pip-audit>=2.7; extra == 'all'
Requires-Dist: pre-commit>=4.2; extra == 'all'
Requires-Dist: pydoclint>=0.5; extra == 'all'
Requires-Dist: pygments<3.0,>=2.20; extra == 'all'
Requires-Dist: pylint>=3.3; extra == 'all'
Requires-Dist: pymdown-extensions>=11.0.1; extra == 'all'
Requires-Dist: pyright>=1.1.403; extra == 'all'
Requires-Dist: pytest-cov>=6.1.1; extra == 'all'
Requires-Dist: pytest-env>=1.0.1; extra == 'all'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'all'
Requires-Dist: pytest-randomly>=3.15; extra == 'all'
Requires-Dist: pytest-timeout>=2.1; extra == 'all'
Requires-Dist: pytest-xdist>=3.6; extra == 'all'
Requires-Dist: pytest>=9.0.3; extra == 'all'
Requires-Dist: pyyaml>=6.0; extra == 'all'
Requires-Dist: radon>=6.0; extra == 'all'
Requires-Dist: ruff>=0.15.2; extra == 'all'
Requires-Dist: toml-sort>=0.23.1; extra == 'all'
Requires-Dist: twine>=6.1; extra == 'all'
Requires-Dist: validate-pyproject>=0.23; extra == 'all'
Requires-Dist: vulture>=2.14; extra == 'all'
Requires-Dist: wandb>=0.18; extra == 'all'
Requires-Dist: xenon>=0.9; extra == 'all'
Provides-Extra: all-cpu
Requires-Dist: bandit[toml]>=1.8.6; extra == 'all-cpu'
Requires-Dist: build>=1.0.3; extra == 'all-cpu'
Requires-Dist: coverage>=7; extra == 'all-cpu'
Requires-Dist: deptry>=0.23; extra == 'all-cpu'
Requires-Dist: griffe>=1.7.3; extra == 'all-cpu'
Requires-Dist: hypothesis>=6.100; extra == 'all-cpu'
Requires-Dist: import-linter>=2.0; extra == 'all-cpu'
Requires-Dist: interrogate>=1.5; extra == 'all-cpu'
Requires-Dist: matplotlib>=3.8; extra == 'all-cpu'
Requires-Dist: mkdocs-material>=9.7.7; extra == 'all-cpu'
Requires-Dist: mkdocs<2.0,>=1.6.1; extra == 'all-cpu'
Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'all-cpu'
Requires-Dist: mkdocstrings>=0.28.3; extra == 'all-cpu'
Requires-Dist: mlflow>=2; extra == 'all-cpu'
Requires-Dist: pip-audit>=2.7; extra == 'all-cpu'
Requires-Dist: pre-commit>=4.2; extra == 'all-cpu'
Requires-Dist: pydoclint>=0.5; extra == 'all-cpu'
Requires-Dist: pygments<3.0,>=2.20; extra == 'all-cpu'
Requires-Dist: pylint>=3.3; extra == 'all-cpu'
Requires-Dist: pymdown-extensions>=11.0.1; extra == 'all-cpu'
Requires-Dist: pyright>=1.1.403; extra == 'all-cpu'
Requires-Dist: pytest-cov>=6.1.1; extra == 'all-cpu'
Requires-Dist: pytest-env>=1.0.1; extra == 'all-cpu'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'all-cpu'
Requires-Dist: pytest-randomly>=3.15; extra == 'all-cpu'
Requires-Dist: pytest-timeout>=2.1; extra == 'all-cpu'
Requires-Dist: pytest-xdist>=3.6; extra == 'all-cpu'
Requires-Dist: pytest>=9.0.3; extra == 'all-cpu'
Requires-Dist: pyyaml>=6.0; extra == 'all-cpu'
Requires-Dist: radon>=6.0; extra == 'all-cpu'
Requires-Dist: ruff>=0.15.2; extra == 'all-cpu'
Requires-Dist: toml-sort>=0.23.1; extra == 'all-cpu'
Requires-Dist: twine>=6.1; extra == 'all-cpu'
Requires-Dist: validate-pyproject>=0.23; extra == 'all-cpu'
Requires-Dist: vulture>=2.14; extra == 'all-cpu'
Requires-Dist: wandb>=0.18; extra == 'all-cpu'
Requires-Dist: xenon>=0.9; extra == 'all-cpu'
Provides-Extra: all-macos
Requires-Dist: bandit[toml]>=1.8.6; extra == 'all-macos'
Requires-Dist: build>=1.0.3; extra == 'all-macos'
Requires-Dist: coverage>=7; extra == 'all-macos'
Requires-Dist: deptry>=0.23; extra == 'all-macos'
Requires-Dist: griffe>=1.7.3; extra == 'all-macos'
Requires-Dist: hypothesis>=6.100; extra == 'all-macos'
Requires-Dist: import-linter>=2.0; extra == 'all-macos'
Requires-Dist: interrogate>=1.5; extra == 'all-macos'
Requires-Dist: jax-metal>=0.1.0; extra == 'all-macos'
Requires-Dist: matplotlib>=3.8; extra == 'all-macos'
Requires-Dist: mkdocs-material>=9.7.7; extra == 'all-macos'
Requires-Dist: mkdocs<2.0,>=1.6.1; extra == 'all-macos'
Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'all-macos'
Requires-Dist: mkdocstrings>=0.28.3; extra == 'all-macos'
Requires-Dist: mlflow>=2; extra == 'all-macos'
Requires-Dist: pip-audit>=2.7; extra == 'all-macos'
Requires-Dist: pre-commit>=4.2; extra == 'all-macos'
Requires-Dist: pydoclint>=0.5; extra == 'all-macos'
Requires-Dist: pygments<3.0,>=2.20; extra == 'all-macos'
Requires-Dist: pylint>=3.3; extra == 'all-macos'
Requires-Dist: pymdown-extensions>=11.0.1; extra == 'all-macos'
Requires-Dist: pyright>=1.1.403; extra == 'all-macos'
Requires-Dist: pytest-cov>=6.1.1; extra == 'all-macos'
Requires-Dist: pytest-env>=1.0.1; extra == 'all-macos'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'all-macos'
Requires-Dist: pytest-randomly>=3.15; extra == 'all-macos'
Requires-Dist: pytest-timeout>=2.1; extra == 'all-macos'
Requires-Dist: pytest-xdist>=3.6; extra == 'all-macos'
Requires-Dist: pytest>=9.0.3; extra == 'all-macos'
Requires-Dist: pyyaml>=6.0; extra == 'all-macos'
Requires-Dist: radon>=6.0; extra == 'all-macos'
Requires-Dist: ruff>=0.15.2; extra == 'all-macos'
Requires-Dist: toml-sort>=0.23.1; extra == 'all-macos'
Requires-Dist: twine>=6.1; extra == 'all-macos'
Requires-Dist: validate-pyproject>=0.23; extra == 'all-macos'
Requires-Dist: vulture>=2.14; extra == 'all-macos'
Requires-Dist: wandb>=0.18; extra == 'all-macos'
Requires-Dist: xenon>=0.9; extra == 'all-macos'
Provides-Extra: cuda-dev
Requires-Dist: bandit[toml]>=1.8.6; extra == 'cuda-dev'
Requires-Dist: build>=1.0.3; extra == 'cuda-dev'
Requires-Dist: coverage>=7; extra == 'cuda-dev'
Requires-Dist: deptry>=0.23; extra == 'cuda-dev'
Requires-Dist: import-linter>=2.0; extra == 'cuda-dev'
Requires-Dist: interrogate>=1.5; extra == 'cuda-dev'
Requires-Dist: jax[cuda12]>=0.11.1; extra == 'cuda-dev'
Requires-Dist: pip-audit>=2.7; extra == 'cuda-dev'
Requires-Dist: pre-commit>=4.2; extra == 'cuda-dev'
Requires-Dist: pydoclint>=0.5; extra == 'cuda-dev'
Requires-Dist: pylint>=3.3; extra == 'cuda-dev'
Requires-Dist: pyright>=1.1.403; extra == 'cuda-dev'
Requires-Dist: pyyaml>=6.0; extra == 'cuda-dev'
Requires-Dist: radon>=6.0; extra == 'cuda-dev'
Requires-Dist: ruff>=0.15.2; extra == 'cuda-dev'
Requires-Dist: toml-sort>=0.23.1; extra == 'cuda-dev'
Requires-Dist: twine>=6.1; extra == 'cuda-dev'
Requires-Dist: validate-pyproject>=0.23; extra == 'cuda-dev'
Requires-Dist: vulture>=2.14; extra == 'cuda-dev'
Requires-Dist: xenon>=0.9; extra == 'cuda-dev'
Provides-Extra: cuda12
Requires-Dist: jax[cuda12]>=0.11.1; extra == 'cuda12'
Provides-Extra: dev
Requires-Dist: bandit[toml]>=1.8.6; extra == 'dev'
Requires-Dist: build>=1.0.3; extra == 'dev'
Requires-Dist: coverage>=7; extra == 'dev'
Requires-Dist: deptry>=0.23; extra == 'dev'
Requires-Dist: import-linter>=2.0; extra == 'dev'
Requires-Dist: interrogate>=1.5; extra == 'dev'
Requires-Dist: pip-audit>=2.7; extra == 'dev'
Requires-Dist: pre-commit>=4.2; extra == 'dev'
Requires-Dist: pydoclint>=0.5; extra == 'dev'
Requires-Dist: pylint>=3.3; extra == 'dev'
Requires-Dist: pyright>=1.1.403; extra == 'dev'
Requires-Dist: pyyaml>=6.0; extra == 'dev'
Requires-Dist: radon>=6.0; extra == 'dev'
Requires-Dist: ruff>=0.15.2; extra == 'dev'
Requires-Dist: toml-sort>=0.23.1; extra == 'dev'
Requires-Dist: twine>=6.1; extra == 'dev'
Requires-Dist: validate-pyproject>=0.23; extra == 'dev'
Requires-Dist: vulture>=2.14; extra == 'dev'
Requires-Dist: xenon>=0.9; extra == 'dev'
Provides-Extra: docs
Requires-Dist: griffe>=1.7.3; extra == 'docs'
Requires-Dist: mkdocs-material>=9.7.7; extra == 'docs'
Requires-Dist: mkdocs<2.0,>=1.6.1; extra == 'docs'
Requires-Dist: mkdocstrings-python>=1.1.2; extra == 'docs'
Requires-Dist: mkdocstrings>=0.28.3; extra == 'docs'
Requires-Dist: pygments<3.0,>=2.20; extra == 'docs'
Requires-Dist: pymdown-extensions>=11.0.1; extra == 'docs'
Provides-Extra: metal
Requires-Dist: jax-metal>=0.1.0; extra == 'metal'
Provides-Extra: mlflow
Requires-Dist: mlflow>=2; extra == 'mlflow'
Provides-Extra: plots
Requires-Dist: matplotlib>=3.8; extra == 'plots'
Provides-Extra: test
Requires-Dist: coverage>=7; extra == 'test'
Requires-Dist: hypothesis>=6.100; extra == 'test'
Requires-Dist: matplotlib>=3.8; extra == 'test'
Requires-Dist: pytest-cov>=6.1.1; extra == 'test'
Requires-Dist: pytest-env>=1.0.1; extra == 'test'
Requires-Dist: pytest-json-report>=1.5.0; extra == 'test'
Requires-Dist: pytest-randomly>=3.15; extra == 'test'
Requires-Dist: pytest-timeout>=2.1; extra == 'test'
Requires-Dist: pytest-xdist>=3.6; extra == 'test'
Requires-Dist: pytest>=9.0.3; extra == 'test'
Provides-Extra: testing
Requires-Dist: pytest>=9.0.3; extra == 'testing'
Provides-Extra: wandb
Requires-Dist: wandb>=0.18; extra == 'wandb'
Description-Content-Type: text/markdown

# Substrax

**JAX/Flax NNX training infrastructure: JAX process configuration, device detection and
placement, device meshes and SPMD sharding (data, FSDP, tensor and pipeline strategies), an
Orbax checkpoint store that restores onto the current devices, early stopping and callbacks,
and W&B/MLflow logging.**
It is the shared layer of the Avitai JAX stack.

[![CI](https://github.com/avitai/substrax/actions/workflows/ci.yml/badge.svg)](https://github.com/avitai/substrax/actions/workflows/ci.yml)
[![Build](https://github.com/avitai/substrax/actions/workflows/build-verification.yml/badge.svg)](https://github.com/avitai/substrax/actions/workflows/build-verification.yml)
[![Docs](https://github.com/avitai/substrax/actions/workflows/docs.yml/badge.svg)](https://github.com/avitai/substrax/actions/workflows/docs.yml)
[![PyPI](https://img.shields.io/pypi/v/substrax.svg)](https://pypi.org/project/substrax/)

[Documentation](https://substrax.readthedocs.io) · [Changelog](CHANGELOG.md) · [Issues](https://github.com/avitai/substrax/issues)

> **Research preview.** Substrax is under rapid iteration and the API will change while
> the sibling packages migrate onto it. Pin a version.

## Where it sits

Substrax is the bottom of the Avitai dependency chain and depends on none of the siblings:

```text
substrax → calibrax → datarax → artifex → opifex
```

It holds the code those packages used to carry separately, so that each concern has one
home and one test suite:

| 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 |
| `substrax.testing` | Opt-in test infrastructure: fresh-interpreter runs with a chosen JAX configuration, and a pytest plugin with `x64`, `devices` and `accelerator` markers and jax configuration isolation |
| `substrax.devices` | `detect_devices()` (platform, device kind, count), 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 and batch placement, `spmd_train_step`, gradient reduction and 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 |

Not in Substrax: optimizer algorithms and schedules (optax, which Substrax assembles from a
config), loss scaling and gradient accumulation
(`flax.training.dynamic_scale.DynamicScale`, `optax.MultiSteps`), profiling and hardware
spec tables (calibrax), data pipelines (datarax), models and trainers (artifex, opifex).

## Installation

```bash
uv add substrax          # or: pip install substrax
uv add "substrax[wandb]"  # Weights & Biases backend
uv add "substrax[mlflow]" # MLflow backend
uv add "substrax[testing]" # pytest plugin and fresh-interpreter test helpers
```

Substrax requires Python 3.12 or 3.13, `jax>=0.11.1`, `flax>=0.12.9` and
`orbax-checkpoint>=0.11.33`. The `cuda12` and `metal` extras select the JAX backend.

## Quick start

One data-parallel step over every visible device, a checkpoint, and an early-stopping
decision. The same code runs on one CPU; on several devices the batch is sharded on its
leading axis and XLA inserts the gradient all-reduce.

```python
import jax
import jax.numpy as jnp
import optax
from flax import nnx

from substrax.callbacks import EarlyStopping
from substrax.checkpoint import OrbaxCheckpointStore
from substrax.devices import detect_devices
from substrax.mesh import DeviceMeshManager
from substrax.spmd import create_data_parallel_sharding, place_batch_on_shards, spmd_train_step

info = detect_devices()  # DeviceInfo(platform='cpu', kind=<DeviceKind.CPU>, count=1, ...)

model = nnx.Linear(8, 1, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)

mesh = DeviceMeshManager.create_device_mesh({"data": info.count})
sharding = create_data_parallel_sharding(mesh)
batch = place_batch_on_shards(
    {"x": jnp.ones((32, 8)), "y": jnp.zeros((32, 1))},
    sharding,
)


def loss_fn(model: nnx.Module, batch: dict[str, jax.Array]) -> jax.Array:
    return jnp.mean((model(batch["x"]) - batch["y"]) ** 2)


stopper = EarlyStopping(patience=3, min_delta=1e-4)
with OrbaxCheckpointStore("checkpoints/quick-start", max_to_keep=2) as store:
    for step in range(5):
        with jax.set_mesh(mesh):
            loss = spmd_train_step(model, optimizer, loss_fn, batch)
        store.save(model, step, float(loss))
        stopper.update(float(loss))  # True when the loss improved on the best so far
        if stopper.should_stop:
            break

    restored, metadata = store.restore(model, store.latest_step())
    assert metadata["loss"] == float(loss)
```

## The subpackages

### Runtime

`JaxRuntime` declares the settings a JAX process starts with: backends, CPU device count,
64-bit types, matmul precision, compilation cache, XLA flags and accelerator memory.
`runtime_environment` renders them for a process that has not imported jax, and
`apply_runtime` applies them to the current one, refusing any setting jax would no longer
read.

```python
import os
import subprocess
import sys

from substrax.runtime import JaxRuntime, runtime_environment

runtime = JaxRuntime(platforms=("cpu",), cpu_devices=8)
env = {**os.environ, **runtime_environment(runtime, os.environ)}
program = "import jax; print(jax.device_count())"
subprocess.run([sys.executable, "-c", program], env=env, check=True)  # prints 8
```

XLA flags merge by name: a flag already set to a different value raises
`XlaFlagConflictError` instead of being replaced. `resolve_test_runtime` picks a test run's
backend and emulated CPU devices, and `configure_entry_point_logging` sets up logging from
`main()` without `force=True`.

### Artifacts

`resolve_output_dir` picks where a run writes its outputs: an explicit directory, then
`$AVITAI_OUTPUT_DIR/<name>`, then a directory created once per process under the system
temporary directory. It never defaults into the working tree, so running an example or a
test cannot overwrite tracked files.

```python
from substrax.artifacts import resolve_output_dir

location = resolve_output_dir("fno_darcy")  # or explicit=Path("docs/assets/examples/fno_darcy")
figure_path = location.path / "prediction.png"
print(location.source)  # "argument", "environment" or "run_default"
```

### Testing

`run_python` runs code in a fresh interpreter with the JAX settings a test chooses, and the opt-in
pytest plugin adds device markers and fails a test that changes jax's global configuration, as
jax's own test harness does.
`TraceCounter` asserts how many times a jitted function traced, for example that a training step
compiles once. `run_example` runs each example in its own interpreter, with its outputs redirected
away from the repository.

```python
# conftest.py
pytest_plugins = ["substrax.testing.pytest_plugin"]

# test_sharding.py
import pytest

from substrax.runtime import JaxRuntime
from substrax.testing import run_python


@pytest.mark.devices(2)
def test_on_two_devices() -> None: ...


def test_eight_emulated_devices() -> None:
    program = "import json, jax; print(json.dumps(jax.device_count()))"
    result = run_python(program, runtime=JaxRuntime(cpu_devices=8), timeout=120)
    assert result.check().last_json() == 8
```

### Devices

`detect_devices()` is the one reading of the hardware the whole stack shares.

```python
from substrax.devices import DeviceKind, detect_devices

info = detect_devices()
assert info.platform in {"cpu", "gpu", "tpu", "metal"}
if info.kind is DeviceKind.GPU:
    print(f"{info.count} GPU(s): {info.device_kinds}")
```

`place_on_device(pytree, device)` moves a pytree, and `get_batch_size_recommendation()`
reads the per-hardware batch-size table.

### Mesh and SPMD

`DeviceMeshManager.create_device_mesh` takes the mesh shape as a mapping from axis name
to size and builds every axis as `Auto` unless `axis_types` says otherwise: with jax 0.11
that is what lets XLA infer the sharding of the backward pass over a batch sharded on the
`data` axis.

```python
import jax
from substrax.mesh import DeviceMeshManager, data_parallel_rules, create_named_sharding
from substrax.spmd import create_data_parallel_sharding

mesh = DeviceMeshManager.create_device_mesh({"data": jax.device_count()})
print(DeviceMeshManager.get_mesh_info(mesh))  # {'total_devices': 1, 'axes': {'data': 1}} on one device

batch_sharding = create_data_parallel_sharding(mesh)      # leading axis over "data"
replicated = create_named_sharding(mesh, None)             # every device holds a copy
rules = data_parallel_rules()                              # MeshRules for nnx.spmd
```

The strategies in `substrax.mesh.strategies` (`DataParallelStrategy`, `FSDPStrategy`,
`TensorParallelStrategy`, `PipelineParallelStrategy`, `MultiDimensionalStrategy`) build
partition specs for a `ParallelismConfig`; `substrax.spmd` adds `reduce_gradient_tree`,
`all_gather` and the `reduce_*` collectives.

### Checkpoint

`OrbaxCheckpointStore` writes a step-addressed store: the payload (an `nnx.Module`, a
`TrainState` or a pytree of arrays) with `PyTreeSave`, and a JSON sidecar with the step,
the loss and any extra metadata. Restoring onto a target places every array on the
target's device, so a checkpoint written on `cuda:0` restores in a CPU-only process.

```python
from flax import nnx
from substrax.checkpoint import OrbaxCheckpointStore

model = nnx.Linear(4, 4, rngs=nnx.Rngs(0))
with OrbaxCheckpointStore("checkpoints/demo", max_to_keep=3) as store:
    store.save(model, step=100, loss=0.25, additional_metadata={"epoch": 2})
    store.save(model, step=200, loss=0.20)

    assert store.list_steps() == [100, 200]
    assert store.best_step("loss") == 200
    fresh = nnx.Linear(4, 4, rngs=nnx.Rngs(1))
    restored, metadata = store.restore(fresh, step=100)  # fresh is updated in place
    assert metadata["epoch"] == 2

    payload, _ = store.restore(step=200)  # no target: the payload as it was stored
```

`restore` with a target raises `ValueError` when the checkpoint's arrays do not fit it,
which is how a store written for another architecture is refused rather than loaded.

### Callbacks

`EarlyStopping` is the plain tracker: `update(value)` records a metric and returns
whether it improved on the best by more than `min_delta`, and `should_stop` is true once
`patience` updates in a row have not. `EarlyStoppingCallback` is the same rule as a
training callback that a trainer drives
through `on_epoch_end`, with `check_finite`, `stopping_threshold` and
`divergence_threshold` from its `EarlyStoppingConfig`.

```python
from substrax.callbacks import EarlyStopping

stopper = EarlyStopping(patience=2, min_delta=0.01, mode="min")
improved = []
for step, loss in enumerate([1.0, 0.5, 0.49, 0.495, 0.5]):
    improved.append(stopper.update(loss))
    if stopper.should_stop:
        break
assert improved == [True, True, False, False]  # 0.49 is not 0.01 better than 0.5
assert step == 3  # two updates in a row without improvement
```

### Tracking

Every logger has the same surface (`log_scalar`, `log_scalars`, `log_hyperparams`,
`log_text`, `log_image`, `log_histogram`, `close`).
`create_logger` gives console output plus a file under `log_dir`; `WandbLogger` and
`MLFlowLogger` need the `wandb` and `mlflow` extras and import their SDK at construction,
so a missing extra fails at the call site, not at import time.

```python
from substrax.tracking import create_logger

logger = create_logger("demo", log_dir="logs")
logger.log_hyperparams({"lr": 1e-3, "batch_size": 32})
for step in range(3):
    logger.log_scalar("loss", 1.0 / (step + 1), step=step)
logger.close()
```

```python
from substrax.tracking import MLFlowLogger, WandbLogger

wandb_logger = WandbLogger("demo", project="my-project", config={"lr": 1e-3})
mlflow_logger = MLFlowLogger("demo", experiment_name="my-experiment")
```

## Development setup

```bash
git clone https://github.com/avitai/substrax.git
cd substrax
./setup.sh
source ./activate.sh
```

`setup.sh` creates the environment with `uv`, syncs the `dev` and `test` extras plus the
backend extra for this machine, and writes the managed environment file `.substrax.env`
that `activate.sh` loads. A user-owned `.env` is never modified.

| Flag | Effect |
| --- | --- |
| `--backend <auto\|cpu\|cuda12\|metal>` | Choose the backend policy; `auto` resolves to `cuda12` on Linux with a visible NVIDIA GPU, `metal` on Apple Silicon, otherwise `cpu` |
| `--python <version>` | Create the environment with a specific Python version |
| `--extra <name>` | Sync an additional extra (repeatable), e.g. `--extra docs` |
| `--recreate` | Remove the existing `.venv` before syncing |
| `--force-clean` | Remove `.venv`, `.substrax.env` and repo-local test artifacts |
| `--dry-run` | Print the resolved backend and the `uv` commands without changing files |

Run the checks the way CI does:

```bash
uv run --locked pytest
uv run --locked pre-commit run --all-files
uv run --locked mkdocs build --strict --clean
```

The test suite also runs as a pre-commit hook, so a commit takes a few seconds longer
than a lint pass.

## Documentation

<https://substrax.readthedocs.io> — one page per subpackage under *API Reference*.

## Contributing

See [CONTRIBUTING.md](CONTRIBUTING.md). Security reports go to the address in
[SECURITY.md](SECURITY.md), not to a public issue.

## License

MIT — see [LICENSE](LICENSE).
