v0.1.0 · MIT · NumPy only

Autodiff from scratch, verified against PyTorch.

TinyTorch implements reverse-mode automatic differentiation, a dynamic computation graph, an nn API and optimizers in about 3,600 lines of Python. NumPy supplies storage and arithmetic — every derivative rule, the graph and the backward traversal are TinyTorch's own.

Read how it works View source

391tests passing

24operators

1runtime dependency

python examples/06_pytorch_equivalence.py

architecture : 6 -> 16 -> 12 -> 4 (tanh)   dtype: float64

[1] forward output
  logits                     max|diff| = 4.441e-16   PASS
[2] loss
  cross-entropy              max|diff| = 4.441e-16   PASS
[3] parameter gradients
  grad 0.weight (16, 6)      max|diff| = 4.163e-17   PASS
  grad 2.weight (12, 16)     max|diff| = 5.551e-17   PASS
  grad 4.weight (4, 12)      max|diff| = 5.551e-17   PASS
[4] after one optimizer step
  SGD(lr=0.1, momentum=0.9)  max|diff| = 5.551e-17   PASS
  Adam(lr=0.01)              max|diff| = 2.220e-16   PASS
[5] after 10 Adam steps
  after step 10              max|diff| = 2.776e-16   PASS

15/15 checks passed at 1e-10

What you actually write

The API is small on purpose. Three things cover almost everything: build a tensor, call backward(), step an optimizer.

Differentiate an expression

Any tensor with requires_grad records the operations applied to it. Calling backward() on a scalar walks that record in reverse and fills in .grad.

Reuse and branching are handled by the traversal, not by you: a tensor feeding two consumers receives the sum of both gradients.

autograd.py

import tinytorch as tt

x = tt.Tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x ** 2 + 3 * x).sum()
y.backward()

x.grad
# array([5., 7., 9.])   ==  2x + 3

Train a network

Module discovers parameters by attribute assignment — writing self.weight = Parameter(...) is the registration. No metaclass, no registry.

zero_grad() is required, not decorative: backward accumulates so that gradient accumulation over micro-batches works.

train.py

from tinytorch import nn, optim

model = nn.Sequential(
    nn.Linear(4, 32), nn.ReLU(),
    nn.Linear(32, 3),
)
opt = optim.Adam(model.parameters(), lr=1e-2)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(100):
    opt.zero_grad()          # backward accumulates
    loss = loss_fn(model(x), y)
    loss.backward()
    opt.step()

Check the gradients yourself

gradcheck compares the analytic gradient against a central finite difference of TinyTorch's own forward pass — no PyTorch involved.

It refuses float32 inputs. The error of a central difference is O(h²) + O(ε/h), and float32's ε leaves no usable window for h. Silently returning garbage would be worse than raising.

gradcheck.py

import numpy as np, tinytorch as tt

x = tt.randn(4, 5, requires_grad=True, dtype=np.float64)
tt.gradcheck(lambda t: tt.log_softmax(t * 2).sum(), [x])
# True

y = tt.randn(4, requires_grad=True, dtype=np.float32)
tt.gradcheck(lambda t: t * 2, [y])
# TypeError: float64 is required

How it is verified

Correctness is established three independent ways. Every number below is measured, and reproducible from the repository.

Against PyTorch, at machine epsilon

84 tests build the same computation in both frameworks from bit-identical inputs and compare in float64 at rtol=1e-10. The completion test builds a 6→16→12→4 tanh MLP with identical weights and measures every stage:

Forward logits4.4e-16

Cross-entropy loss4.4e-16

Parameter gradients (all 6)≤ 1.1e-16

After one SGD / momentum / Adam step≤ 2.2e-16

After 10 Adam steps2.8e-16

five orders of magnitude inside the enforced tolerance

Against finite differences, and against reality

101 tests check every operator numerically — framework-independent, so they would catch an error TinyTorch and PyTorch happened to share. A deliberately wrong derivative is included in the suite to prove the checker catches one.

17 more train real models and assert against known answers: recovered generating parameters, the least-squares optimum, a held-out accuracy floor. A control test confirms a linear model fails XOR — so the MLP's success is attributable to the nonlinearity, not to an easy dataset.

Operators vs finite differences101

PyTorch parity84

Module, tensor, optimizer, serialization189

Real training convergence17

Total, all passing391

Passing with NumPy alone307

CI: Python 3.10–3.13 · Linux, macOS, Windows

Overhead, measured

TinyTorch is not trying to be fast. The benchmark exists so the cost of a Python autograd engine is a measured number rather than a guess. CPU, float64, best of five repeats, torch.set_num_threads(1).

Workload	TinyTorch	PyTorch	Ratio
`add`, 1×1	0.0062 ms	0.0095 ms	0.7×
`matmul`, 32×32	0.0227 ms	0.0239 ms	1.0×
`matmul`, 256×256	0.355 ms	0.341 ms	1.0×
`exp`, 256×256	0.313 ms	0.160 ms	2.0×
Training step, 698 params	0.133 ms	0.112 ms	1.2×
Training step, 2.1 M params	37.6 ms	24.2 ms	1.6×
1000-op chain, fwd + bwd	6.20 ms	2.62 ms	2.4×

The result contradicted the prediction

The expected finding was the standard one: fixed Python overhead amortises away as models grow. The measurements show the opposite.

Faster at 1×1Not a win worth claiming — PyTorch's full C++ dispatcher costs more per call than a Python path supporting exactly one backend. Both sit in the same microsecond band.
Parity where BLAS runsmatmul is 1.0× at every size tested. Both frameworks call essentially the same kernel.
Slower at scaleThe training-step ratio rises from 1.2× to 1.6×. What TinyTorch loses at scale is kernel efficiency and extra temporaries, not dispatch.
Linear in graph depth~6 µs per op, flat across 10, 100 and 1000 ops — no accidental quadratic behaviour in the traversal.
A correctionThe first version of this benchmark timed tensor construction along with dispatch. Leaves are now built outside the timed region; the numbers above are from the corrected run.

Scope is frozen

TinyTorch exists to demonstrate how autodiff, tensors, neural-network abstractions and optimizers work, in code short enough to read end to end. Everything below is absent by design — not a roadmap, not a TODO. Use PyTorch for real work.

No double backwardGradients are ndarray, not Tensor, so the backward pass records no graph. Rules out gradient penalties and MAML-style meta-learning.
CPU onlyNo CUDA, MPS or ROCm, and no device abstraction. There is no .to(device) and there is not meant to be one.
No convolutions or recurrenceNo Conv2d, BatchNorm, LayerNorm, Embedding, LSTM or attention. The op set covers dense networks.
No training infrastructureNo schedulers, no DataLoader, no distributed training, no mixed precision, no profiler.
No in-place mutationThere is no version counter, so mutating a tensor a saved context still references would silently corrupt its gradient.
Memory is not managedEvery saved tensor stays alive for the whole backward pass. No checkpointing, no rematerialisation.

Read the internals

Both documents are written to be read start to finish, with diagrams and derivations.

Document

Reverse-mode autodiff

Why reverse mode and not forward, how the dynamic graph is recorded, the chain rule as implemented, topological traversal, gradient accumulation, and broadcast gradient reduction.

Read → Document

Architecture

Layering and the import cycle, Module/Parameter discovery by attribute interception, optimizer state, initialization, and the pickle-free checkpoint format.

Read →

install

$ git clone https://github.com/Gariyuuu/tinytorch && cd tinytorch
$ python -m venv .venv && source .venv/bin/activate
$ pip install -e ".[dev]"
$ pytest
391 passed