MyTorch

A deep learning framework you can read end to end, built to learn how PyTorch works.

MyTorch is a from-scratch automatic differentiation and neural network library, written for educational purposes. It started as a reimplementation of micrograd and grew into something that follows PyTorch's API closely enough that what you learn here transfers, while being small enough to actually read: about 10,600 lines, every gradient rule derived by hand and visible in the file next to the operation it belongs to.

It is not a PyTorch replacement and does not try to be one. PyTorch is faster, vastly more capable, and battle-tested by an enormous number of people. If you are training a real model, use PyTorch. Use this if you want to see how one works.

from mytorch import Value

x, y = Value(-2.0), Value(3.0)
z = (x * y + y.tanh()).sigmoid()
z.backward()

x.grad  # 0.019847...

---

What this is for

Reading a real framework's source is hard because production code is shaped by performance, backwards compatibility, and hardware support — PyTorch's ATen is hundreds of thousands of lines across C++, Python and generated code. The ideas underneath it are much smaller than that, and this is an attempt to show them at a size a person can hold in their head:

• How reverse-mode autodiff actually works — scalar.py is the whole idea in one file, one node per number, no dependencies at all. Read it first; tensor.py is the same design at array scale.
• How nn.Module finds your parameters — attribute assignment, in nn.py, including why a tied weight must be returned once.
• How optimizers really differ — SGD, Adam and AdamW in optim.py, with tests that pin the exact arithmetic of the first two steps rather than just checking the loss fell.
• What a compiler pass over a graph looks like — compile.py lowers a graph to a flat instruction stream and does liveness analysis to reuse buffers.
• Forward mode, and why both modes exist — forward.py, cross-checked against reverse mode to 1e-12.

Every claim in this README has a test or a benchmark behind it, and where a number came out badly it is printed anyway.

Training diagnostics that name the problem

This is the one part of MyTorch that does something PyTorch has no equivalent for, so it gets its own section.

When a network trains badly, the loss curve tells you that and nothing more. The cause is usually one of a small number of specific conditions inside the network, and all of them are visible in numbers you can measure — if you know which numbers, and what values are bad. PyTorch ships every raw ingredient for this: register_forward_hook, .grad, and TensorBoard's add_histogram. What it does not ship — and what the third-party dashboards built on it do not ship either, because they plot what you give them — is anything that reads those numbers and tells you what is wrong.

diagnostics.py does that:

from mytorch.diagnostics import TrainingMonitor

# Used as a context manager, so the forward hooks it installs are always removed.
with TrainingMonitor(model) as monitor:
    for step in range(steps):
        loss = loss_fn(model(x), y)
        optimizer.zero_grad()
        loss.backward()
        # After backward and before the step: the one moment when the activations
        # and their gradients both exist.
        monitor.record(learning_rate=lr)
        optimizer.step()

print(monitor.report())

Real output, two findings of ten, wrapped to fit:

diagnosis: 10 finding(s)
  [saturated-activations] layers.1 (Tanh): 82% of units have |activation| > 0.97, where
  the derivative is under 0.06, so little gradient passes back through this layer
  (reported above 66%)
    try: scale down the initialization of the layer feeding this one, or put an
         nn.BatchNorm1d before it so its input arrives normalized

  [learning-rate-too-high] layers.0.bias: log10(update/data) is -1.15, so each step
  moves this parameter by 7.0% of its own scale (healthy is near -3, meaning 0.1%)
    try: lower the learning rate

It recognizes saturated activations, vanishing and exploding gradients, dead ReLU units, NaN or infinite values, a gradient that never arrives, and a learning rate too high or too low — each per layer or per parameter, so you get a place to look and not just a verdict. Every message quotes the threshold it applied, because the thresholds are conventions rather than laws and you should be able to disagree with a specific number instead of with the tool.

The four recorded views are the ones Andrej Karpathy works through in makemore Part 3: forward activation statistics, backward gradient statistics, per-parameter grad:data, and the update-to-data ratio over time — log10(|lr·grad| / |data|), which should sit near −3. plot_activations, plot_gradients and plot_update_ratios draw them, with the −3 reference line on the last. The statistics themselves need nothing beyond NumPy; only the plots want matplotlib.

examples/diagnose_training.py reproduces the video's experiment on one architecture, changing only the initialization:

$ python examples/diagnose_training.py --compare
setup       saturation  grad in/out  loss                 diagnosis
-------------------------------------------------------------------
saturated   82%         2.66e+00     41.948 -> 3.537      saturated-activations
vanishing   0%          2.15e-07     0.872 -> 0.871       vanishing-gradients
batchnorm   4%          9.27e-02     21.820 -> 3.199      none
healthy     13%         8.79e-02     1.202 -> 0.768       none

The vanishing row is the case worth reading twice: over 20 steps the loss moves by 0.001, and the diagnosis says why — the gradient reaching the first layer is 2e-7 times the size of the gradient at the last. That is the symptom you cannot see in a loss curve, which is the whole argument for measuring it. --setup batchnorm keeps the bad initialization and adds nn.BatchNorm1d; both faults go away, which is the video's conclusion, reproduced here as a test rather than a claim.

What this is not. It is a set of well-chosen measurements and a rule-based reading of them, not a trained model or a novel result. Each rule is a threshold an experienced practitioner already carries in their head. The value is that a beginner does not have to, and that a practitioner does not have to re-derive the plots for every project.

Small, deliberate departures from PyTorch

These are not improvements on PyTorch's engineering. They are places where a teaching framework can afford a stricter choice that a production one cannot, because MyTorch has no backwards-compatibility obligations and no users to break.

1. Broadcasting is strict by default. This is a genuine semantic difference, not a nicer error message. Forget keepdim in a reduction and NumPy aligns a (3,) vector to the wrong axis: on a non-square tensor you get an error, but on a square one you get the wrong number and a model that trains almost well enough to look fine. MyTorch refuses the ambiguous case:

>>> square = Tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
>>> square - Tensor([1., 2., 3.])
BroadcastError: cannot subtract shapes (3, 3) and (3,) under strict broadcasting.
  left:  shape (3, 3) float64 from train.py:2 in <module>
  right: shape (3,) float64 from train.py:4 in <module>
  Ranks differ (2 vs 1), so the alignment is ambiguous: NumPy would align the (3,)
  operand to the trailing axes of (3, 3), which is a silent wrong answer if you
  meant rows.
  Fix: reshape it to (1, 3) to broadcast across rows, or (3, 1) across columns.
  Or opt into NumPy rules: mytorch.set_broadcast_policy("numpy").

NumPy's rules remain available with one call, for porting existing code.

2. Undefined operations raise instead of returning nan. torch.log(torch.tensor(-1.)) returns nan and lets it spread through a training run; here it raises, naming the value and the fix. This is mostly better error messages, and it is a real convenience when you are learning — a nan that appears six layers away from its cause teaches you nothing.

3. Every tensor records the line of your code that created it, so an error points at your source rather than at an anonymous operation inside the library.

4. Physical units, which no mainstream framework carries. units.py propagates dimensions through the graph, so kg + g and exp(hours) are refused. Genuinely useful for scientific work, and a small enough idea to read in one sitting.

What is implemented

Two engines sharing one design: a scalar teaching kernel and a NumPy-backed tensor engine. The API deliberately mirrors torch, so this doubles as a map of PyTorch's own surface.

Tensor engine	Elementwise arithmetic, exp/log/sqrt/tanh/sigmoid/relu/abs, matmul, reductions over any axis with keepdims, reshape, transpose, indexing, broadcast_to — every one with a hand-derived VJP
Strict broadcasting	Mismatched ranks raise instead of silently aligning to the wrong axis. NumPy rules remain available for porting
Provenance	Every tensor records the line of your code that made it, and errors quote it
Scalar engine	The same API one number at a time, with zero dependencies — the version you can read in a sitting
Networks	mytorch.nn — Module with parameter auto-discovery, Linear (with lazy shape inference), activations, BatchNorm1d, Sequential, Flatten, MSELoss, CrossEntropyLoss, state_dict, forward hooks
Training diagnostics	mytorch.diagnostics — records the four makemore-Part-3 views and names the fault: saturation, vanishing/exploding gradients, dead units, NaN, learning rate too high or low, per layer. See above
Optimizers	mytorch.optim — SGD (momentum, Nesterov, weight decay), Adam, AdamW
Data	mytorch.data — Dataset, TensorDataset, and a seeded in-process DataLoader
Native kernels	mytorch.backend — an optional Rust extension, 2.1-3.2x faster on a fused sum-of-squares, with the kernels that lose to NumPy reported too
Second derivatives	mytorch.hvp gives exact Hessian-vector products by reverse-over-forward, checked against closed-form Hessians
Declared invariants	mytorch.check_invariants tests invariance, equivariance, conservation and bounds by sampling, reporting the worst deviation and the input that caused it
Forward-mode AD	mytorch.jvp for Jacobian-vector products, verified against reverse mode to 1e-12
Neural ODEs	mytorch.odeint backpropagates through an RK4 solver; odeint_adjoint integrates backwards for memory constant in the step count
Physical units	mytorch.units propagates dimensions through the graph and refuses exp(hours) or kg + g, quoting the line of your code that built the operand
Memory autopilot	mytorch.plan_segments solves checkpoint placement exactly under a byte budget, and refuses up front if none fits
Compiled tape	mytorch.lower flattens a graph into instructions over pooled buffers: 2-3x faster replay than eager, and 10-30x less intermediate memory
Run recording	mytorch.Run logs config, seed, git commit and loss curve, and records a crash instead of losing the run
Scalar networks	mytorch.scalar_nn — Neuron, Layer, MLP over Value
Verification	mytorch.gradcheck, plus 1201 Python tests and 5 Rust tests including finite-difference checks on every operation and closed-form pins on every optimizer update
Determinism	One seeded generator; nothing touches NumPy's global random state
Runtime dependencies	One: NumPy. The scalar engine needs nothing at all

Performance, stated carefully

Every speedup below compares MyTorch against itself. None of it is a comparison against PyTorch — PyTorch appears in this project only as a correctness oracle, and no speed comparison against it was ever run. Vectorizing is worth what you would expect, and the reason is visible in the graph size rather than the arithmetic:

Network, batch	Scalar engine	Tensor engine	Speedup	Scalar nodes	Tensor nodes
2-16-16-1, batch 8	56.2 ms	4.7 ms	12x	6,218	41
2-16-16-1, batch 256	4495.4 ms	4.1 ms	1095x	167,418	41
2-64-64-1, batch 64	18789.4 ms	4.3 ms	4356x	570,954	41

Forward plus backward, minimum of 5 repetitions, Python 3.11.15 and NumPy 2.4.6 on an Intel i5-10300H. The tensor engine's node count does not depend on the batch at all, which is the whole story. This compares MyTorch's two engines against each other, not against PyTorch — see benchmarks/ for what the numbers do and do not mean, and reproduce them with uv run python benchmarks/scalar_vs_tensor.py.

That last row is deliberate, and measurable. The wheel is 118 KB; the whole installed stack including NumPy is about 20 MB, against roughly 2.5 GB for a CUDA-bundled PyTorch wheel (pytorch#17621). NumPy is the tensor engine's one dependency, and it is the only one: CI fails the build if the wheel exceeds 1 MB, if installing it pulls in anything beyond NumPy, or if mytorch.scalar — the teaching kernel — stops importing in an environment with nothing installed at all. That last gate is why the NumPy-backed names are imported lazily.

Install

pip install mytorch-ml

The import name is mytorch; the distribution name carries an -ml suffix because mytorch was claimed on PyPI by an unrelated 2018 placeholder.

Optional extras: mytorch-ml[viz] for Graphviz rendering, mytorch-ml[examples] to run the demos.

Quickstart

Train a network the way you would expect to:

import numpy as np
import mytorch
from mytorch import Tensor, data, nn, optim

mytorch.manual_seed(0)

model = nn.Sequential(nn.Linear(32), nn.ReLU(), nn.Linear(2))  # no input width needed
loader = data.DataLoader(data.TensorDataset(features, labels), batch_size=32, shuffle=True)
loss_fn = nn.CrossEntropyLoss()

model(Tensor(features[:1]))  # lazy layers size themselves here
optimizer = optim.Adam(model.parameters(), lr=0.01)

for epoch in range(5):
    for batch_x, batch_y in loader:
        loss = loss_fn(model(batch_x), batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Wrap it in a Run and the seed, git commit, config and loss curve are recorded for you — see examples/digits_classifier.py, which reaches ~99% test accuracy on scikit-learn's digits in 8 epochs.

Train a classifier on two interleaving half-moons:

uv run --extra examples python examples/moons_demo.py --steps 100 --seed 42

Verify a gradient you are unsure about:

from mytorch import gradcheck

gradcheck(lambda a, b: (a * b).tanh() / b.log(), (0.7, 1.3))  # True, or raises

Look at the graph you actually built:

from mytorch import Value
from mytorch.viz import render_graph

loss = ((Value(2.0) * Value(-3.0)) + Value(10.0)).relu()
loss.backward()
render_graph(loss, "artifacts/graph")  # -> artifacts/graph.svg

How far it goes, and where it stops

The project was built in ten releases, each tagged with a record of what was verified. It is complete enough to train real (small) models and to read as a reference. Below is what exists and — more usefully for a learner — where each piece stops.

The engines. A scalar engine (Value, one node per number, zero dependencies) and a NumPy tensor engine sharing the same design, so you can read the small one first. Both are complete for the operations listed above; neither has convolutions, GPU support, or float32.

A compiler pass. mytorch.compile lowers a graph into a flat instruction stream over pooled buffers, with liveness analysis so buffers are reused: 3.2x faster replay than our own eager mode at batch 1 falling to 2.0x at batch 256, intermediate memory down 10-30x. The trend is the interesting part — the tape removes a fixed per-operation Python cost, so it helps most where that cost dominates the arithmetic.

Native kernels. crates/mytorch-core is a small Rust extension (PyO3) reached through mytorch.backend, optional, with a NumPy fallback. Measured against NumPy, and the losses are published alongside the wins because that is the honest shape of the result:

kernel	vs NumPy
sum_of_squares	2.1x to 3.2x faster
a * b + c fused	1.1x to 1.35x faster
tanh(a * b) + c fused	0.87x, slower
matmul	0.04x, 25x slower than OpenBLAS

Fusion pays where arithmetic is cheap and the cost is memory traffic. It loses to NumPy's SIMD-vectorized transcendentals, and a readable cache-blocked triple loop is nowhere near a hand-tuned BLAS microkernel. So the two losing kernels default to native=False with their measurements in their docstrings, and benchmarks/native_vs_numpy.py prints both columns.

Memory for compute. plan_segments solves gradient-checkpoint placement under a byte budget exactly rather than heuristically — the objective mixes a sum over stored boundaries with a maximum over segment sizes, so candidate maxima are enumerated and a dynamic program runs under each as a cap. Verified against exhaustive search over every partition. (Rotor is the lineage.)

Both modes of autodiff, and second derivatives. jvp for forward mode, verified against reverse mode to a relative 1e-12; hvp for exact Hessian-vector products by reverse-over-forward.

Science contracts. Physical units with dimensional analysis; declared invariants (invariance, equivariance, conservation, bounds) checked by sampling; neural ODEs with the adjoint method. The pharmacokinetics example learns a saturable drug-elimination law and records three findings about fitting mechanisms that were kept rather than tuned away.

Deliberately not built: convolutions, GPU support, distributed training, float32, vmap, a graph rewriter. Any of those would be a reasonable exercise on top of what is here.