Get Started

pip install tensorax

from tensorax import Tensor, nn, optim, lr_scheduler, functional as F

# Build
model = nn.Sequential(nn.Linear(4, 8), nn.GELU(), nn.LayerNorm(8), nn.Linear(8, 3))
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

# Train
for epoch in range(100):
    loss = F.mse_loss(model(x_train), y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    scheduler.step()

→ Full usage guide with all APIs, code examples, and details: docs/USAGE.md

What's Inside

Performance

Matrix Multiplication — fp32, 3×1024×1024, 100 runs:

PyTorch CUDA (ref)         ████████████████████████████████████████████  0.41s  (4.51×)
Tensorax 1D Block Tiling   ██████████████████████████████████████████    0.95s  (2.31×)  ← best
Tensorax Tiled             ████████████████████████████████              1.22s  (1.80×)
NumPy CPU (baseline)       █████████████████████████                    1.85s  (1.00×)

2.31× faster than NumPy · 43% of PyTorch's cuBLAS kernels · all hand-written, zero library calls

Attention Kernels — fp32/fp16, B=4, H=8, S=256, Dk=512, Dv=512, 30 runs:

PyTorch SDPA (ref)         ████████████████████████████████████████████  0.04s  (2340×)
Tensorax MMA Tensor Core   ██████████████████████████████████████        0.33s   (297×)  ← best
Tensorax Optim. Flash      ██████████████████████████████████            0.52s   (187×)
Tensorax Flash SDPA        ██████████████████████████                    3.10s    (31×)
NumPy CPU (baseline)       ████████████████████                          7.06s    (14×)
Tensorax Tiled SDPA        ██████                                       32.91s     (3×)
Tensorax Naive SDPA        █                                            98.26s     (1×)

9.3× faster than Flash SDPA via raw PTX inline assembly using mma.sync Ampere Tensor Cores and SFU intrinsics.

Project Structure

csrc/                           C++ / CUDA backend
  cuda/kernels/                   elementwise · matmul (×6) · reduction · attention (×5)
  cpu/                            CPU fallback for all ops
  tensor_ops.{cpp,h}             pybind11 bindings

tensorax/                       Python package
  tensor.py                       Tensor class + autograd
  functional.py                   F.relu, F.gelu, F.silu, F.softmax, F.sdpa, ...
  nn/                             Linear, Embedding, norms, dropout, attention (SDPA, MHA, GQA)
  optim.py                        SGD, Adam
  lr_scheduler.py                 StepLR, CosineAnnealingLR, ExponentialLR, LinearLR, MultiStepLR

Roadmap

Documentation

Contributing

Fork → Branch → Commit → PR

See DEVELOPMENT.md for build instructions and guidelines.

@software{tensorax2025,
  title  = {Tensorax: Pure C++/CUDA Tensor Library},
  author = {Shrirang Mahajan},
  year   = {2025},
  url    = {https://github.com/NotShrirang/tensorax}
}