Metadata-Version: 2.4
Name: raxis
Version: 0.1.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: dev
Requires-Dist: raxis[build]; extra == "dev"
Requires-Dist: raxis[test]; extra == "dev"
Requires-Dist: raxis[examples]; 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. CNN layers and CUDA are 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`.

## 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, NEON, ...) and picked at runtime |
| `libraxis/raxis/autograd` | Graph, engine and backward functions |
| `libraxis/bindings` | pybind11 module `raxis._C` |
| `raxis/` | Python package: `nn`, `optim`, 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, and
`examples/` has notebooks that train MLPs on two moons and MNIST with either
raxis or PyTorch.

## 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).
