Axis

crates.io · latest GitHub release · Why Axis · Vision · Library contract · Contributing · License · Trademarks

Deep-learning programs are unusually good at being wrong while continuing to run. A loader wraps around, and a fresh-data study silently becomes a repeated one. A test example slips into training under a new seed or a new file name. A class axis is averaged away by accident. The loss still falls, the job still exits zero, and the number that gets reported describes a different experiment from the one on paper. No error is raised, because nothing was ever checked.

Axis is a Rust framework for training neural networks on NVIDIA cuTile that checks. The assumptions a result depends on are executable: when one stops being true, the run stops before the bad step reaches the optimizer, and when it holds, the run hands back a receipt that says exactly what was verified. Named tensor axes do the same for the math, so a batch axis can never quietly become a class axis.

The same experiment in every framework. Each table is one matched run: byte-identical data, sample order and initial weights, the same optimizer and float32 throughout, each framework on its default training path. The losing numbers stay in.

ImageNet64, a three-block CNN, one full epoch (1,281,167 images), top-1 and top-5 on all 50,000 held-out images, on one RTX 5070 Ti:

	Top-1	Top-5	Images/s, steady	Training time	Lines of code	Catches train/eval overlap	Catches data reuse
Axis 0.11	0.0540	0.1485	782	1,647 s	173	rejects the run	rejects the run
PyTorch 2.11	0.0554	0.1511	30,269	186 s	101	not checked	not checked
Lightning 2.6	0.0552	0.1504	26,120	183 s	116	not checked	not checked
JAX 0.11	0.0548	0.1499	31,559	45 s	111	not checked	not checked

Fashion-MNIST, a small MLP, five passes, median of three runs, on one RTX 5060 (another session shared that GPU during timing, and the receipt says so):

	Accuracy	Training time, total	Training time, steady	Lines of code	Catches train/eval overlap	Catches data reuse
Axis 0.11	0.5918	7.41 s	0.035 s	155	rejects the run	rejects the run
PyTorch 2.11	0.5918	0.17 s	0.016 s	73	not checked	not checked
Lightning 2.6	0.5918	0.20 s	0.071 s	77	not checked	not checked
JAX 0.11	0.5918	2.94 s	0.042 s	82	not checked	not checked
Burn (CUDA)	0.5938	2.86 s	0.082 s	118	not checked	not checked
Candle (CUDA)	0.5918	0.03 s	0.018 s	114	not checked	not checked

Every framework learns the same thing: the arms start with identical accuracy and end within 0.0014 top-1 (ImageNet64) and 0.002 (Fashion-MNIST) of each other. Axis is the slowest end to end in both runs, and at convolutional scale it is about 39 times slower than PyTorch; on the small MLP its steady step is faster than Lightning, JAX and Burn, and slower than PyTorch and Candle. Axis also takes the most code, and that count includes the checks. It is the only framework that stops when an evaluation sample leaks into training or a sample repeats within a declared pass. These are bounded runs on single hosts, not framework-wide claims. Receipts: ImageNet64 and Fashion-MNIST, from axis-benchmarks 480e7ee.

Project status: Axis is active experimental research software. It is CUDA-only, and its API can change as new consumer programs expose better boundaries. The repository records measured witnesses, but does not claim production readiness or framework-wide performance parity.

Add the current crates.io release to a Rust project:

cargo add [email protected]

Executable research checks

These are what Axis is for. Each one guards an assumption a result depends on, runs inside the training loop, fails before the offending step reaches the optimizer, and returns a receipt that says exactly what was verified.

Assumption	Check	What it catches
Every sample is new	SinglePass, FinitePasses	a step-count edit that turns a fresh-sample run into repeated passes over a finite corpus
Repetition stays within a declared rate	Idr with IdrLimits	a generator or loader that repeats examples more than the experiment allows
Train and evaluation never share an example	Disjointness over an IdentityScheme	overlap by meaning, not by file or seed: the same puzzle after relabeling, the same document through two shards
The model actually learned	LearningProgress	a loop that runs every forward, backward and step while one declared metric on one declared evaluation population never improves
A property holds on ordered inputs	EmpiricalMonotonicity	violations of a claimed monotone relation, recorded as sampled evidence rather than a global guarantee
A physical law is satisfied	EmpiricalResidual	a model whose residual against a named law, region and evaluator exceeds its declared limit
The run can be reproduced	Trainer::step_training with a TrainingPass	nondeterministic randomness: every random layer draws from a seed derived from the run seed and the step, and the step receipt names it

Axis distinguishes data claims that ordinary loaders collapse into one:

single pass                 no finite example is intentionally reused
finite passes               reuse is explicit and counted
no exact sample reuse       observed stable identities do not repeat
IDR approximation           declared coverage and repeat limits still hold
semantic disjointness       canonical identities do not cross populations

A receipt never claims more than was checked. Passing assert_idr() does not prove independent samples or a rich underlying distribution; it proves the declared operational conditions. A sampled monotonicity check says "sampled", and learning progress says "improved by this much on this population", not "converged". Receipts (IdrReceipt, DisjointnessReceipt, LearningProgressReceipt, ...) are the fragments a run certificate is built from. The contracts are in data regimes, semantic disjointness, learning progress, static guarantees and the library design.

A training loop

use axis::prelude::*;

let (batch, input, hidden, output) = (
    Axis::new("batch"),
    Axis::new("input"),
    Axis::new("hidden"),
    Axis::new("output"),
);

let mut model = Sequential::new((
    Linear::new(input, hidden.of(16)),
    ReLU,
    Linear::new(hidden, output.of(1)),
));
model.build(&Shape::new([batch.of(256), input.of(2)])?, &device, 42)?;

let mut data = DataLoader::new(AdditionDataset::new(0xadd1_7100), 256)?
    .assert_idr(IdrLimits::generated(0.0)?)?;
let mut trainer = Trainer::new(SGD::new(0.25)?);

There are no epochs in this example because the source never ends. If a draw ID repeats, IDR fails. A separate Disjointness ledger compares canonical problem identities across training and evaluation, so the run also fails when distinct draw IDs encode the same reference problem. Streaming mode discards training keys after comparison, keeping memory bounded by the retained evaluation data.

Sudoku as an acceptance test

Axis includes a Rust migration of furkanhaney/sudoku-transformer. It is a useful framework test because the task is easy to inspect while the program exercises a real transformer: learned token and position embeddings, bidirectional multi-head attention, pre-LayerNorm residual blocks, GELU, blank-only categorical loss, AdamW, and whole-puzzle accuracy.

The data source generates valid boards and fresh clue masks indefinitely. The same run tests Axis's infinite-data-regime assertion and a semantic disjointness ledger whose identity is the clue board after canonical digit relabeling.

flowchart LR
    A[Valid generated boards] --> B[Fresh clue masks]
    B --> C[DataLoader + assert_idr]
    C --> D[Token + position embeddings]
    D --> E[Pre-norm bidirectional attention]
    E --> F[GELU feed-forward block]
    F --> G[9 logits per cell]
    G --> H[Loss on blanks only]
    H --> I[AdamW update]
    C -. canonical puzzle identity .-> J[Train/tuning/audit disjointness]

Loading

The first bounded run trained a 3,129-parameter acceptance model on 50 fresh boards. Evaluation loss fell from 2.8734 to 2.2069; all 50 training IDs were unique and overlap with the eight fixed evaluation boards was zero. Accuracy remained near chance, so this is a mechanics result rather than a Sudoku-solving claim.

On a rented RTX 5090, the same tiny model completed 200 updates over 1,600 fresh boards: evaluation loss moved from 2.8734 to 2.1707, blank accuracy reached 14.93%, reuse and evaluation overlap remained zero, and exact solves remained zero. Peak GPU utilization was only 6% for that run, which makes launch count and contraction lowering the measured performance frontier rather than a throughput success claim.

Conformance with PyTorch

Axis is not a PyTorch clone, but a model ported from PyTorch should compute the same thing, and that is checked rather than assumed.

• Feature parity. PyTorch 2.14's torch.nn layer catalog is frozen as a 161-row spec and each class is graded against Axis's own definition of a finished module: a named-axis contract, explicit state, an independent oracle for forward and every gradient, CUDA coverage including reordered storage, and documented defaults. Today 112 yes, 16 partial, 33 no: 72.05% (catalog). scripts/checks/nn_gap.sh --check recomputes the number on every gate and fails if it drops below its floor; the rows still no are the ones the backlog defers on purpose (lazy modules, data-parallel wrappers, fractional pooling) and the ones waiting on persistent training state (BatchNorm).
• Numerical conformance. Every operator is tested forward and backward against an independent oracle written as literals or reimplemented from PyTorch's documented formula, never computed by the code under test. The full gate runs 226 CUDA tests and 67 host tests.
• Against PyTorch itself. pytorch_activation_forward_and_gradient_parity runs PyTorch alongside Axis (set AXIS_PYTHON to an interpreter with torch) and compares values and gradients. Eight research studies ported from PyTorch (a byte language model, a Voronoi bottleneck, an upscaler, a cell partitioner, a multiple-instance classifier, a Fourier MLP, an energy fit and a U-Net block) each carry a small Axis crate that reproduces a PyTorch forward and loss, with observed error below 1e-5 in every case.
• Reference architectures. axis::architectures ships ResNet-34, ResNet-50, VGG16 (with and without BatchNorm) and DCGAN, built only from public modules. Each matches its reference's parameter count exactly (ResNet-50: 25,557,032, the same as torchvision) and, at a reduced width loaded with the reference's weights, its logits and first-layer gradient to within 1e-7.
• Known, stated divergences are named on the operator, for example gelu is PyTorch's tanh form and gelu_exact its default, and a logsumexp group of all non-finite values returns NaN where PyTorch returns negative infinity.

Execution

The backend enqueues each eager training step on one CUDA stream and synchronizes once at the step boundary. Single-axis contractions use batched tiled GEMM; multi-axis contractions retain the generic gather/reduce path. Layout permutations and split/merge use rank-sized device indexing, retaining identity transforms as storage views. Broadcast/reduction operations reuse CPU and bounded device-side index plans. FP32 is the default, with an explicit BF16-matrix/FP32-state device policy. Convolution materializes FP32 patch tensors before contraction; volumetric patches can therefore dominate memory, and later generic broadcast/reduction operations retain their own index-plan limits. Set AXIS_PROFILE=1 to print Trainer phases, tensor planning, submission, read, and synchronization timings while investigating a workload. These are host wall times, not CUDA kernel durations; see the execution profile.

Learning paths and measured programs

Programs are grouped by the background needed to understand their claim. This is an audience ladder, not a leaderboard: verification, monotonicity, and data-regime contracts do not sit beside introductory training merely because all of them are executable.

Framework verification stays with its owner in src/library/ tests, repository tests/, and data/evidence/. Empirical monotonicity and IDR belong to research-contract programs or outside research consumers. They are checks on a declared claim, not another kind of beginner example.

Getting started

Program	Purpose	Current witness
MNIST	recognizable finite-pass categorical training	bounded smoke improves held-out accuracy from 10.16% to 33.98%
Fashion-MNIST	train-fitted preprocessing and a less separable grayscale task	bounded smoke improves held-out accuracy from 9.77% to 47.66%
CIFAR-100	named-axis RGB convolution and 100-class training	bounded smoke lowers pass loss to 4.4359 and improves held-out accuracy to 2.54%
ImageNet64	licensed large-data preparation and a compact 1,000-class CNN	deterministic parser tests plus a real-pixel one-step GPU receipt; no accuracy claim

Training building blocks

These are ordinary model and optimizer flows. “Training” is broader and more accurate than “layers”: Muon is an optimizer, while CNN, MLP, and attention are compositions.

Program	Purpose	Current witness
MLP	library and explicit cuTile baseline	both reach 5.74e-2 held-out MSE after 500 steps
CNN	convolution forward/backward and learning	configured grouped convolution matches a scalar oracle; a depthwise-separable block composes; smoke reaches 100% accuracy
Attention	causal attention and prefix-mean learning	central differences pass; held-out MSE falls to 2.39e-2
Muon	explicit Muon/AdamW partition and paired learning	both identically initialized arms reduce held-out MSE by more than 99% in 100 steps; no winner claim

Research contracts

Program	Purpose	Current witness
Generated addition	endless data with executable IDR assumptions	128,000 fresh samples, zero observed reuse; held-out MSE falls to 3.88e-3
MNIST population	independent models and rates on one population axis	fused four-member smoke improves best accuracy from 8.20% to 60.55%
Sudoku transformer	bidirectional transformer plus generated IDR acceptance	50 unique training boards, zero evaluation overlap, finite forward/backward/update
Chess transformer	geometric attention, joint policy/value learning, finite game-disjoint data	two AdamW updates, exact pass receipt, zero train/evaluation overlap

Applied studies

Program	Purpose	Current witness
Country panel	AdamW and split/preprocessing migration	bounded validation MSE falls on country and future splits; no GDP-fit claim
Damped pendulum	physics-informed learning from a nonlinear ODE	RK4 angle/velocity RMSE reach 0.00749 rad and 0.01102 rad/s; strict sampled residual limit passes

These numbers are repository witnesses with different tasks and budgets. They show that the exercised path works; they are not a benchmark leaderboard.

Project structure

axis/
├── src/
│   ├── library/              published Axis framework crate
│   ├── examples/
│   │   ├── getting-started/
│   │   │   ├── cifar100/      100-class RGB convolution
│   │   │   ├── fashion-mnist/ clothing classification
│   │   │   ├── imagenet64/    licensed large-data CNN starting point
│   │   │   ├── mnist/         approachable end-to-end training
│   │   │   └── vision-data/   strict shared dataset and training adapters
│   │   └── training/
│   │       ├── attention/    causal-attention consumer and oracle
│   │       ├── cnn/          CNN consumer and scalar oracle
│   │       ├── mlp/          MLP consumer and explicit cuTile baseline
│   │       └── muon/         paired Muon and AdamW learning acceptance
│   └── studies/
│       ├── contracts/
│       │   ├── generated-addition/  generated-data IDR training
│       │   └── mnist-population/    population-axis research contract
│       ├── energy-output/
│       │   └── panel/        country-year regression migration
│       └── physics/
│           └── damped-pendulum/ nonlinear ODE residual acceptance
├── docs/                     shared contracts, assumptions, and next work
├── data/                     tracked evidence and selected runs; other contents ignored
├── scripts/                  CUDA setup, Cargo runner, checks, and censuses
├── Cargo.toml                workspace and shared dependency versions
└── Cargo.lock                one resolved dependency graph

The original Python programs behind the migrations remain in the sibling research repository. Each migration document names the source and separates preserved mechanics from claims it has not reproduced.

Setup

Requirements:

• Linux and an NVIDIA GPU supported by cuTile
• Rust 1.89 or newer
• libclang
• CUDA 13.2

Provision the repository-local CUDA toolkit when the host does not already have a compatible toolkit:

bash scripts/setup_cuda.sh

The generated build/cuda/ and Cargo target/ directories stay local. CUDA's six library aliases are materialized as independent regular files after every setup invocation, including when all package markers already match. The setup rejects any other symlink in the local toolkit. The only external direct Rust dependency is pinned cutile = "=0.3.1".

Run

Run commands from the repository root:

bash scripts/check.sh

bash src/examples/getting-started/mnist/scripts/train.sh --smoke
bash src/examples/training/mlp/scripts/train.sh --smoke
bash src/examples/training/cnn/scripts/train.sh --smoke
bash src/examples/training/attention/scripts/train.sh --smoke
bash src/examples/training/muon/scripts/train.sh --smoke
bash src/studies/contracts/generated-addition/scripts/train.sh --smoke
bash src/studies/contracts/mnist-population/scripts/train.sh --smoke
bash src/studies/energy-output/panel/scripts/train.sh --smoke --split country
bash src/studies/physics/damped-pendulum/scripts/train.sh --smoke

check.sh runs formatting, Clippy, CPU tests, and serialized CUDA tests. The ignored test suite needs an NVIDIA device and is intentionally driven by the script.

Where Axis is going

The project thesis is to give coordinates to the whole experiment, not only its tensors. It grows against an acceptance ladder: the Sudoku and Chess transformers are current, and running the sub-30B model behind ask_second_opinion is the long-range systems test.

The current implementation frontier is recorded in next work.

Contributing and license

Human and agent-assisted contributions are welcome. Start with the contribution guide, which explains how changes are scoped, verified, and presented for review.

Original Axis material is available under the MIT License. Dependencies, downloaded toolchains and datasets, and linked external artifacts retain their own licenses and terms.