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 and close to it on NVIDIA GPUs (without cuDNN).

Status: alpha. CPU and CUDA. 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), on the CPU or on an NVIDIA GPU.

Install

pip install raxis

Wheels for Windows (x86-64), Linux (x86-64 and aarch64, e.g. a 64-bit Raspberry Pi OS) and macOS (Apple silicon, macOS 14+), Python 3.10 to 3.14. Raxis has no runtime dependencies; numpy is optional (.numpy(), raxis.from_numpy).

For NVIDIA GPUs, install it with the CUDA backend (Windows and Linux x86-64):

pip install "raxis[cuda12]"   # driver r525+, GPUs from Volta (sm_70) on
pip install "raxis[cuda13]"   # driver r580+, GPUs from Turing (sm_75) on

The backend is a package of its own, as JAX's CUDA plugins, so the CPU wheel stays small: pip install raxis-cuda12 / raxis-cuda13 is the same. It pulls the raxis of its version and NVIDIA's cuBLAS wheel; the CUDA runtime is linked in, so only the driver is needed. import raxis loads it, and raxis.cuda.is_available() tells whether there is a GPU to use.

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 my two earlier projects: 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's Conv2d is SynapGrad's algorithm on libtorch (sliding windows and tensordot), not PyTorch's convolution. SynapX supports up to PyTorch 2.7, so its column runs in an environment of its own.

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.

On the GPU. RTX 5070 Ti (CUDA 13.0, driver 610), float32, PyTorch 2.14.1+cu130, the same TF32 flags in both libraries (nanoGPT turns TF32 on for matmuls, as its train.py does). Median training step:

The full nanoGPT run takes 6.0 min with Raxis and 5.8 min with PyTorch, best val loss 1.4668 and 1.4612 (both at iteration 1750). PyTorch convolves with cuDNN; Raxis does not link it (a closed library of several hundred MB that would hide the most interesting kernels) and runs im2col plus cuBLAS. Its attention is its own FlashAttention-2 style kernel, against PyTorch's memory-efficient one. At batch 64 the MLP is bound by kernel launches on any GPU.

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.

ARM. On GitHub's ARM runners (Apple M1 and Neoverse, shared machines, so no table) the median ratio Raxis/PyTorch over the same cases is 0.79x to 0.88x in float32 and float64.

Against NumPy. Raxis wheels are 8-10 MB, smaller than NumPy's (13-17 MB), while PyTorch's CPU wheels take 110-180 MB. On 1000 x 1000 elementwise ops Raxis and PyTorch both run 3 to 30x faster than NumPy (add 0.017 vs 0.565 ms, exp 0.123 vs 0.944 ms), mostly because they split the work between threads: NumPy's own functions are single-threaded by design, only its BLAS calls use several. On 100 x 100 arrays the three are on par.

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; with --cuda it also builds the CUDA backend (needs the CUDA toolkit, 12.8 or newer, which brings cuBLAS), and import raxis loads that build before any installed raxis-cuda* 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, scripts/lineage measures what made my earlier projects slow, 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, Intel MKL (x86-64), OpenBLAS (Linux aarch64) and, in the CUDA packages, the CUDA runtime and CUB: see THIRD_PARTY_NOTICES.txt.