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.
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file substrax-0.1.0.tar.gz.
File metadata
- Download URL: substrax-0.1.0.tar.gz
- Upload date:
- Size: 37.7 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
9b8ec6166c0202e11fd92339e7e2e9baed2b05b49aeb5fd8e45d3517be1b3c83
|
|
| MD5 |
42983866847cd3912097faf4da223b15
|
|
| BLAKE2b-256 |
30bbcf33bc10a454c5b2250a6e570f5ade4ba7ce6e62fd70328b249d795ffdb3
|
Provenance
The following attestation bundles were made for substrax-0.1.0.tar.gz:
Publisher:
publish.yml on avitai/substrax
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
substrax-0.1.0.tar.gz -
Subject digest:
9b8ec6166c0202e11fd92339e7e2e9baed2b05b49aeb5fd8e45d3517be1b3c83 - Sigstore transparency entry: 2763584074
- Sigstore integration time:
-
Permalink:
avitai/substrax@c2b7141006c4c6db09fd2d9ee726d2b15fb2e5aa -
Branch / Tag:
refs/heads/main - Owner: https://github.com/avitai
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@c2b7141006c4c6db09fd2d9ee726d2b15fb2e5aa -
Trigger Event:
workflow_dispatch
-
Statement type:
File details
Details for the file substrax-0.1.0-py3-none-any.whl.
File metadata
- Download URL: substrax-0.1.0-py3-none-any.whl
- Upload date:
- Size: 44.8 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
fba35f0d6844f46c41f2cdfaa8f4d9eccff3e0477ab4cb8bae1b99e3c38640b1
|
|
| MD5 |
c90c74636417a4de1149ffcadaa5d1f9
|
|
| BLAKE2b-256 |
cc336fcfa61962aa8e05e2a758ca41f09246c56c4b742372caa9fc6676cae6bd
|
Provenance
The following attestation bundles were made for substrax-0.1.0-py3-none-any.whl:
Publisher:
publish.yml on avitai/substrax
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
substrax-0.1.0-py3-none-any.whl -
Subject digest:
fba35f0d6844f46c41f2cdfaa8f4d9eccff3e0477ab4cb8bae1b99e3c38640b1 - Sigstore transparency entry: 2763584134
- Sigstore integration time:
-
Permalink:
avitai/substrax@c2b7141006c4c6db09fd2d9ee726d2b15fb2e5aa -
Branch / Tag:
refs/heads/main - Owner: https://github.com/avitai
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@c2b7141006c4c6db09fd2d9ee726d2b15fb2e5aa -
Trigger Event:
workflow_dispatch
-
Statement type:
