Raxis

Raxis is a from-scratch re-implementation of the core of PyTorch, the ATen tensor library and the autograd engine, in C++20 with Python bindings. It is distilled from the PyTorch source: it keeps the techniques that make ATen fast (TensorIterator, vectorized kernels with runtime CPU dispatch, BLAS/MKL) without the breadth and indirection that make the original hard to read.

The goal is a codebase that someone curious about how PyTorch works inside can read end to end, while running at PyTorch speed on CPU.

Status: alpha. CPU only. Tensors with views and broadcasting, type promotion, elementwise ops, reductions, matmul, autograd, and an nn/optim package with enough to train MLPs, CNNs (Conv1d/2d, pooling, BatchNorm) and a small GPT (Embedding, LayerNorm, GELU, scaled dot product attention). CUDA is next.

Install

pip install raxis

Wheels for Windows and Linux (x86-64), Python 3.10 to 3.14. Raxis has no runtime dependencies; numpy is optional (.numpy(), raxis.from_numpy).

import raxis
from raxis import nn, optim

model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

x, y = raxis.randn((64, 784)), raxis.randint(0, 10, (64,))
loss = nn.functional.cross_entropy(model(x), y)
optimizer.zero_grad()
loss.backward()
optimizer.step()

The API follows PyTorch's, so most code only needs torch replaced by raxis.

raxis.autograd.graph exports the autograd graph of a tensor as Graphviz DOT or Mermaid text, as autograd nodes (view="backward") or as forward ops (view="forward"):

from raxis.autograd import graph

print(graph.to_mermaid(loss, dict(model.named_parameters()), view="forward"))
graph.render(loss, "loss.svg")  # needs pip install 'raxis[viz]' and Graphviz

Performance

Measured on a Ryzen 7 9700X (Zen 5, 8 cores, AVX512) with Windows 11, 8 threads, PyTorch 2.14.1 and Intel MKL in both libraries. The scripts and notebooks to reproduce them are in the repo (see Build from source).

Training the examples. Time of trainer.fit in each notebook (training plus the validation pass of every epoch) with Raxis, PyTorch and the two earlier projects of the same author: synapx (autograd in C++ on top of libtorch) and synapgrad (autograd in pure Python on top of NumPy).

The models reach the same accuracy with every library (MNIST: ~97.6% val with the MLP, ~98.6% with the CNN). Most of PyTorch's time in the CNN is Dropout: its CPU random numbers come from a serial mt19937, Raxis draws them from Philox in parallel. synapx has no Conv2d.

nanoGPT. nanoGPT's character-level Shakespeare model, trained with examples/nanogpt/train.py (nanoGPT's model.py with only the imports changed):

Median time of a training step (forward, backward, gradient clipping and AdamW). The full config would take ~5.1 h with Raxis and ~7.2 h with PyTorch on this CPU for its 5000 iterations, but the val loss stops improving around iteration 2000 and the model overfits afterwards, so both were stopped early: the table has the best val loss up to iteration 2000 (1750 for Raxis, 2000 for PyTorch), against 1.4697 in nanoGPT's README. Random numbers differ between the libraries, so the losses do too.

Ops. A selection of scripts/bench_ops.py (float32, forward only unless noted). Over its 388 cases the median ratio Raxis/PyTorch is 0.85x; the only case beyond 2x is the large 3x3 convolution, where PyTorch uses oneDNN's direct convolution and Raxis im2col plus an MKL gemm.

Inside

Code taken from PyTorch keeps a // From: aten/src/ATen/... comment pointing to the original.

Build from source

Requirements: a C++20 compiler (MSVC 17+, GCC 12+, Clang 14+), CMake 3.22+ and uv.

git clone --recursive https://github.com/pgmesa/raxis.git
cd raxis
uv sync --extra dev --extra cpu        # or --extra cu130 for CUDA PyTorch
uv run python scripts/build_libraxis.py
uv run pytest tests

uv sync installs the pinned build and test dependencies from uv.lock, including Intel MKL from PyPI and PyTorch, the reference the tests compare against. build_libraxis.py builds libraxis with CMake and installs it into the raxis/ package. To work on the C++ side, libraxis/CMakePresets.json has debug and release presets; CMakeLists.txt lists the build options (RAXIS_USE_*).

scripts/bench_ops.py benchmarks each op against PyTorch and NumPy, scripts/bench_training.py times a training step of the example models, and examples/ has notebooks that train MLPs on two moons and MNIST and a CNN on MNIST with either Raxis or PyTorch, and examples/nanogpt trains nanoGPT's character-level Shakespeare model, with its model.py unchanged except for the imports, and samples from it. raxis.save and raxis.load use PyTorch's checkpoint format, so a model trained with one library can be sampled with the other.

License

MIT. Raxis includes code derived from PyTorch (BSD-3-Clause) and links SLEEF, mimalloc and Intel MKL: see THIRD_PARTY_NOTICES.txt.