Metadata-Version: 2.4
Name: raxis
Version: 0.2.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: 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.

> **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

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

```python
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"`):

```python
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](#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](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.20 s | 1.4 s | 0.53 s | 0.26 s | 337 | 150 | 4 | 50 |
| mnist_mlp.ipynb | 2.9 s | 4.5 s | 5.6 s | 10.0 s | 178,110 | 48,000 | 64 | 10 |
| mnist_cnn.ipynb | 10.8 s | 37.8 s | - | 4 min 55 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 has no Conv2d.

**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.

**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 |

## 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/autograd` | Graph, engine and backward functions |
| `libraxis/bindings` | pybind11 module `raxis._C` |
| `raxis/` | Python package: `nn`, `optim`, `autograd.graph`, tensor printing |

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. 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](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 and Intel MKL: see [THIRD_PARTY_NOTICES.txt](THIRD_PARTY_NOTICES.txt).
