Substrax

Shared training and hardware infrastructure for the Avitai JAX stack.

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

Substrax is the bottom of the Avitai dependency chain:

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.devices`	Device information and device placement, including the batch-size recommendation table
`substrax.mesh`	Device meshes on `jax.make_mesh`, mesh rules and partition-spec helpers, sharding strategies on `flax.nnx.spmd`
`substrax.spmd`	Data-parallel placement, `spmd_train_step`, gradient reduction and collectives on the flat-state path
`substrax.checkpoint`	One `CheckpointStore` protocol and one Orbax implementation over `CheckpointManager`
`substrax.callbacks`	The training-callback protocol, `BestMetricTracker`, `EarlyStopping`
`substrax.tracking`	Step-wise experiment tracking with console, CSV, Weights & Biases and MLflow backends

Not in Substrax: optimizers and schedules (optax), 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

uv add substrax   # or: pip install substrax

Substrax requires Python 3.12 or later, jax>=0.11.1, flax>=0.12.9 and orbax-checkpoint>=0.11.33.

Development setup

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:

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

Documentation

https://substrax.readthedocs.io

Contributing

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

License

MIT — see LICENSE.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

substrax-0.1.2.tar.gz (38.0 kB view details)

Uploaded Sep 9, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

substrax-0.1.2-py3-none-any.whl (45.1 kB view details)

Uploaded Sep 9, 2026 Python 3

File details

Details for the file substrax-0.1.2.tar.gz.

File metadata

Download URL: substrax-0.1.2.tar.gz
Upload date: Sep 9, 2026
Size: 38.0 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for substrax-0.1.2.tar.gz
Algorithm	Hash digest
SHA256	`0b549852dc1e1519ac72ebb5294aae03c2645f3e206e0f6673c4d1f0c7f2462e`
MD5	`9a771e34cb94d7e75bf3d285a1486a2e`
BLAKE2b-256	`82353c52d3b96127345646eaf6bc3b9d9571ca40f9b73c86b0fd1e5fb71cea9b`

See more details on using hashes here.

Provenance

The following attestation bundles were made for substrax-0.1.2.tar.gz:

Publisher: publish.yml on avitai/substrax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: substrax-0.1.2.tar.gz
- Subject digest: 0b549852dc1e1519ac72ebb5294aae03c2645f3e206e0f6673c4d1f0c7f2462e
- Sigstore transparency entry: 2763780292
- Sigstore integration time: Sep 9, 2026
Source repository:
- Permalink: avitai/substrax@48167c4229b200efcd988ecf24f3db4175a6519e
- Branch / Tag: refs/heads/main
- Owner: https://github.com/avitai
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: publish.yml@48167c4229b200efcd988ecf24f3db4175a6519e
- Trigger Event: workflow_dispatch