substrax 0.1.0__tar.gz

substrax-0.1.0/.gitignore

@@ -0,0 +1,267 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+# generated by setup.sh
+.substrax.env
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer,
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Editor ignore files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+
+# Temp & Cache
+temp/
+tmp/
+*.tmp
+cache/
+
+# JAX / XLA
+.cache/jax/
+.cache/xla/
+.jax_cache/
+jax_cache/
+
+# ML Artifacts
+*.npy
+*.npz
+*.hdf5
+*.h5
+checkpoints/
+*.orbax-checkpoint-tmp-*/
+
+# Experiment tracking
+wandb/
+.wandb*
+mlruns/
+.mlflow/
+tensorboard_logs/
+
+# Benchmark results & data
+benchmark_results/
+benchmark-data/
+.benchmarks/
+
+# Logs
+logs/
+*.log
+
+# OS
+.DS_Store
+Thumbs.db
+
+# IDE (additional)
+.idea/
+.vscode/
+
+# Generated by setup.sh
+
+# Internal / Private
+internal_docs/
+memory-bank/
+documents/
+sandbox/
+CLAUDE.md
+.claude/
+.cursor/
+.agent/
+.taskmaster/
+
+# Private working notes and per-editor assistant configuration. Local only.
+AGENTS.md

substrax-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+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.

substrax-0.1.0/PKG-INFO

@@ -0,0 +1,322 @@
+Metadata-Version: 2.5
+Name: substrax
+Version: 0.1.0
+Summary: Substrax: shared training and hardware infrastructure for the Avitai JAX stack
+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.5.0,>=1.24
+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: optax>=0.2.8; 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: optax>=0.2.8; 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: optax>=0.2.8; 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: optax>=0.2.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: wandb
+Requires-Dist: wandb>=0.18; extra == 'wandb'
+Description-Content-Type: text/markdown
+
+# 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:
+
+```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.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
+
+```bash
+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
+
+```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
+```
+
+## Documentation
+
+<https://substrax.readthedocs.io>
+
+## 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).

substrax-0.1.0/README.md

@@ -0,0 +1,80 @@
+# 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:
+
+```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.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
+
+```bash
+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
+
+```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
+```
+
+## Documentation
+
+<https://substrax.readthedocs.io>
+
+## 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).