Skip to content
Back to work

TensorTrail

A neural network framework built on NumPy, from the autograd up.

State
Shipped
Category
AI / ML
Year
2026
Role
Solo
Deployment
Local onlyPython package, installed locally
TensorTrail gradient flow
  1. Tensor

    NumPy array that remembers its origin

  2. Operation graph

    built at runtime, node per op

  3. Topological sort

    reverse order for the backward pass

  4. Autograd

    local derivative per node

  5. Broadcast reduction

    gradients summed back to shape

  6. Optimiser step

    SGD, momentum, or Adam

TensorTrail is a neural network framework written on top of NumPy: a reverse-mode autograd engine, layers, losses, optimisers, a trainer, serialisation, and a gradient checker. It exists because I could use PyTorch fluently and could not have told you, precisely, what happens when you call backward().

Why build one

There is a difference between knowing that gradients flow backwards through a graph and being able to write the code that makes them. The first is a sentence. The second requires deciding how a tensor remembers where it came from, how a broadcast is undone on the way back, and what a topological order means when the graph is built at runtime.

TensorTrail is small enough to read end to end and complete enough to train MLPs and small convolutional networks. It is not competing with PyTorch. It is the version of PyTorch you can hold in your head.

The autograd engine

The core is a NumPy-backed Tensor that records the operation that produced it. Every operation appends a node to a dynamic graph, and the backward pass walks that graph in topological order, accumulating gradients as it goes.

Broadcasting is where the subtlety lives. NumPy will happily expand a (1, n) tensor against an (m, n) one on the forward pass, which means the backward pass has to sum the gradient back down to the original shape. Getting that wrong produces a model that trains, slowly and incorrectly, which is the worst possible failure mode because nothing crashes.

Graphs can be exported to Graphviz DOT, which turns "why is this gradient zero" from a print-statement exercise into a picture.

What is in it

  • Layers - Linear, Conv2D, pooling, normalisation, dropout, and activations, composable through Sequential.
  • Losses - mean squared error, binary cross entropy, BCE with logits for numerical stability, and multi-class cross entropy.
  • Optimisers - SGD, SGD with momentum, and Adam, with global gradient clipping and epoch-level StepLR and ExponentialLR schedules.
  • Data - offline synthetic datasets and a mini-batch DataLoader with drop_last.
  • Trainer - validation, metrics, early stopping, named checkpoints, and .npz save and load.
  • Gradient checker - centred finite differences, which is the only thing that will tell you your backward pass is wrong before your loss curve does.

Testing

An autograd engine with a subtle bug still trains. That makes tests less optional here than almost anywhere else, and the suite is correspondingly large for the size of the codebase: thirteen test modules covering the tensor, operations, autograd, modules, losses, metrics, optimisers, data loading, the CNN path, serialisation, graph export, the trainer, and the gradient checker itself.

python -m pytest -q
........................................................................ [ 52%]
................................................................. [100%]

137 passed in 2.13s

That run is from a clean clone of the repository, on 21 August 2026. CI runs the same suite on every push through GitHub Actions.

What it changed

Mostly it changed how I read other people's machine-learning code. Framework behaviour that used to be a black box, gradient accumulation, why in-place operations are dangerous, why the stable loss and the naive loss are different functions, is now something I have implemented and debugged. It also made me considerably more suspicious of a loss curve that goes down.

Have an interesting problem?

I am most useful on the ones that cross a boundary: a product decision that turns out to be a data-model decision, or a bug that only reproduces in production. If that sounds like yours, I would like to hear about it.

vanshajpoonia06@gmail.com