Metadata-Version: 2.4
Name: raxis
Version: 0.4.0
Summary: A readable re-implementation of the core of PyTorch (ATen tensor library + autograd) in C++20, with Python bindings.
Author-email: Pablo García Mesa <pgmesa.sm@gmail.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/pgmesa/raxis
Keywords: Tensor,Autograd,Deep Learning,PyTorch,ATen,C++,SIMD
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Education
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: Microsoft :: Windows
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: C++
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Software Development
Classifier: Topic :: Software Development :: Libraries
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: <3.15,>=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES.txt
Provides-Extra: build
Requires-Dist: pybind11<3.0.0,>=2.13.0; extra == "build"
Requires-Dist: pybind11-stubgen<3.0.0,>=2.5.0; extra == "build"
Requires-Dist: numpy<3.0.0,>=2.0.0; extra == "build"
Requires-Dist: mkl-devel>=2025.0; (platform_machine == "x86_64" or platform_machine == "AMD64") and extra == "build"
Requires-Dist: mkl-static>=2025.0; (platform_machine == "x86_64" or platform_machine == "AMD64") and extra == "build"
Requires-Dist: intel-openmp>=2025.0; (platform_machine == "x86_64" or platform_machine == "AMD64") and extra == "build"
Provides-Extra: test
Requires-Dist: pytest<9.0.0,>=8.0.0; extra == "test"
Requires-Dist: numpy<3.0.0,>=2.0.0; extra == "test"
Provides-Extra: examples
Requires-Dist: numpy<3.0.0,>=2.0.0; extra == "examples"
Requires-Dist: matplotlib>=3.8.0; extra == "examples"
Requires-Dist: scikit-learn>=1.4.0; extra == "examples"
Requires-Dist: tqdm>=4.66.0; extra == "examples"
Requires-Dist: ipywidgets>=8.0.0; extra == "examples"
Requires-Dist: ipykernel>=6.0.0; extra == "examples"
Provides-Extra: viz
Requires-Dist: graphviz>=0.20; extra == "viz"
Provides-Extra: cuda12
Requires-Dist: raxis-cuda12; extra == "cuda12"
Provides-Extra: cuda13
Requires-Dist: raxis-cuda13; extra == "cuda13"
Provides-Extra: dev
Requires-Dist: raxis[build]; extra == "dev"
Requires-Dist: raxis[test]; extra == "dev"
Requires-Dist: raxis[examples]; extra == "dev"
Requires-Dist: raxis[viz]; extra == "dev"
Provides-Extra: cpu
Requires-Dist: torch<3.0.0,>=2.0.0; extra == "cpu"
Requires-Dist: torchvision>=0.15.0; extra == "cpu"
Provides-Extra: cu130
Requires-Dist: torch<3.0.0,>=2.0.0; extra == "cu130"
Requires-Dist: torchvision>=0.15.0; extra == "cu130"
Dynamic: license-file

# 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

```bash
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):

```bash
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](#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](https://github.com/pgmesa/synapx) (autograd in C++ on top of libtorch)
and [SynapGrad](https://github.com/pgmesa/synapgrad) (autograd in pure Python
on top of NumPy).

| Notebook | Raxis | PyTorch | SynapX | SynapGrad | Params | Train set | Batch | Epochs |
|---|--:|--:|--:|--:|--:|--:|--:|--:|
| basic_mlp.ipynb | 0.22 s | 0.50 s | 2.0 s | 0.33 s | 337 | 150 | 4 | 50 |
| mnist_mlp.ipynb | 3.0 s | 4.6 s | 6.3 s | 9.4 s | 178,110 | 48,000 | 64 | 10 |
| mnist_cnn.ipynb | 12.5 s | 47.2 s | 3 min 7 s | 5 min 6 s | 20,538 | 48,000 | 128 | 5 |

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](https://github.com/karpathy/nanoGPT)'s
character-level Shakespeare model, trained with `examples/nanogpt/train.py`
(nanoGPT's `model.py` with only the imports changed):

| Config | Raxis step | PyTorch step | Raxis val loss | PyTorch val loss |
|---|--:|--:|--:|--:|
| nanoGPT README's CPU run: 4 layers, n_embd 128, block 64, batch 12, 2000 iters | 13.2 ms | 14.2 ms | 1.9065 | 1.9121 |
| `train_shakespeare_char.py`: 6 layers, n_embd 384, block 256, batch 64, dropout 0.2 | 2.61 s | 4.00 s | 1.4553 | 1.4680 |

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:

| Model | Raxis | PyTorch |
|---|--:|--:|
| MNIST MLP, batch 64 (`scripts/bench_training.py --device cuda`) | 0.54 ms | 0.45 ms |
| MNIST CNN, batch 128 (same) | 1.32 ms | 0.96 ms |
| nanoGPT `train_shakespeare_char.py`, 5000 iters | 47.3 ms | 45.7 ms |

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.

| Op | Shape | Raxis (ms) | PyTorch (ms) |
|---|---|--:|--:|
| add | 1000 x 1000 | 0.017 | 0.017 |
| exp | 1000 x 1000 | 0.123 | 0.130 |
| sum(dim=1) | 1000 x 1000 | 0.010 | 0.018 |
| mm | 512 x 512 @ 512 x 512 | 0.459 | 0.464 |
| softmax | 16384 x 384 | 0.872 | 0.885 |
| layer_norm | 16384 x 384 | 0.726 | 1.182 |
| gelu | 16384 x 384 | 0.814 | 0.800 |
| sdpa (causal) | 64 x 6 x 256 x 64 | 13.98 | 15.23 |
| sort | 1000 x 1000 | 5.98 | 11.03 |
| x[idx] backward | 1000 x 1000 | 0.050 | 0.317 |
| normal_ | 1000 x 1000 | 0.548 | 3.111 |
| dropout | 1000 x 1000 | 0.123 | 6.880 |
| conv2d 3x3, 32 -> 64 ch | 64 x 32 x 26 x 26 | 2.98 | 1.49 |
| conv2d 3x3, 64 -> 64 ch | 64 x 64 x 32 x 32 | 8.94 | 4.32 |

**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](https://numpy.org/doc/stable/reference/global_state.html),
only its BLAS calls use several. On 100 x 100 arrays the three are on par.

## Inside

| Path | What it holds |
|---|---|
| `libraxis/raxis/tensor` | TensorImpl and storage, views and strides, type promotion, TensorIterator |
| `libraxis/raxis/ops` | Operators and their CPU kernels, compiled once per ISA (AVX2, AVX512, NEON) and picked at runtime |
| `libraxis/raxis/ops/cuda` | The CUDA kernels: elementwise loops, reductions, attention, convolutions, sort, indexing |
| `libraxis/raxis/cuda` | CUDA runtime: caching allocator, pinned host memory, cuBLAS handles |
| `libraxis/raxis/autograd` | Graph, engine and backward functions |
| `libraxis/bindings` | pybind11 module `raxis._C` |
| `raxis/` | Python package: `nn`, `optim`, `cuda`, `autograd.graph`, tensor printing |
| `plugins/cuda` | The `raxis-cuda12` / `raxis-cuda13` packages |

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](https://docs.astral.sh/uv/).

```bash
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](https://github.com/karpathy/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](THIRD_PARTY_NOTICES.txt).
