A machine learning library written from scratch in Python, with a deep learning compiler that generates C.
No PyTorch, no TensorFlow, no NumPy in the core. Every operation, from the autograd engine to the attention mechanism, is implemented from first principles. The compiler on top of it analyses computation graphs, fuses operations, and emits optimised C at run time.
A differentiable tensor library. Reverse-mode automatic differentiation on a define-by-run graph, a typed memory backend, and broadcasting that follows NumPy semantics without depending on NumPy.
A neural network framework. A Module system with automatic parameter registration, layers, optimisers, and loss functions.
A decoder-only transformer. Multi-head causal self-attention, LayerNorm, GELU, learned positional encoding, and pre-norm residual blocks, all built on the autograd engine above and gradient-checked individually.
Hardware backends. Matrix multiplication dispatches to the Apple Metal GPU, to Apple Accelerate BLAS on the CPU, or to a pure Python fallback, chosen by problem size.
A deep learning compiler. A separate front end that builds a graph, plans operator fusion, generates C source, compiles it with gcc, and loads it back through ctypes.
Monte Carlo option pricing. Derivatives priced on the tensor library, including the random number generation, and checked against the closed-form answers where those exist.
A 550,977-parameter GPT trained on the 1.11M-character tinyshakespeare corpus, using data-parallel training across 8 CPU cores.
| Model | 4 layers, 4 heads, 128 dimensions |
| Training | 33,113 steps, roughly 8 hours |
| Loss | 3.66 to 1.44 (character-level cross-entropy) |
The model learns Shakespeare's structure and register from raw characters, with no hand-coded rules:
ROMEO:
You will make him to her hell fame:
Be false and by the news of him.
Thou hast not late with my loath oCld,
NORTHUMBERLAND:
To the likely faint with my name are whose inecius:
I cannot know your grace; I and my heart!
It learned the style, not the meaning. Words wander and lines do not resolve, which is the honest ceiling for a model this size. Real coherence needs roughly two orders of magnitude more parameters.
Watch it learn
Samples taken during the run, showing what the model had worked out at each point. Nothing about play structure was ever specified.
Step 1,000. Letter frequencies and word lengths. No words.
Nos sor nore and thee. Kacefrar, sand a air to a trimeng ow brage;
Be cives'd suee ie fie icires eater,
Step 2,000. It discovers that plays have named speakers, and invents one.
OLIZABETH:
The heave in withat's mairest shis shalll weere pooosinessour centery
Step 3,000. Reaching for a real character name, and nearly getting there.
CORIONUS:
What you sumbled?
Step 33,113. Real speakers, verse structure, archaic register.
ROMEO:
You will make him to her hell fame:
Be false and by the news of him.
A call option pays the amount a share finishes above an agreed level
The payoff is the rectifier. The same function that keeps gradients alive in a neural network is the one that determines what a derivative pays, so the whole pricing routine runs as a chain of tensor operations the library already had.
Under the risk-neutral measure
and the share is modelled by geometric Brownian motion,
Applying Itô's lemma to
The
The expectation is then estimated by simulation. Draw
By the strong law of large numbers
so a hundredfold increase in paths buys roughly a tenfold reduction in error, and the convergence table below shows exactly that. Note the rate is independent of dimension, which is why simulation remains viable for payoffs where deterministic quadrature does not.
Normal draws come from the Box-Muller transform, since the library has no dependency that provides them. From two independent uniforms
are independent standard normals. The construction is the polar form of the Gaussian:
Variance is reduced further by antithetic variates, pairing each draw with its reflection:
Since
Six views of 200,000 simulated futures. Red paths finish above the strike and pay out, grey ones expire worthless. The fan shows the range of outcomes widening with time, which is why an option costs anything at all. The gold bars in the last panel are the payouts themselves.
Priced with a share at 100, a strike of 100, 20% volatility, 5% interest, one year to expiry:
| Option | Simulated | Exact |
|---|---|---|
| Ordinary call | 10.4525 | 10.4506 |
| Asian, ordinary average | 5.8696 | no formula exists |
| Asian, geometric average | 5.6532 | 5.6411 |
| Barrier, knocked out at 130 | 3.8233 | no formula used |
The ordinary call has the Black-Scholes closed form
with
The exotics. These depend on the whole path
and the up-and-out barrier payoff is
Only the geometric Asian admits a formula, and the reason is structural: a product of lognormals is lognormal, whereas a sum of them is not. That single case is enough to validate the path simulation, after which the other two are trusted.
The measured prices obey the inequalities they must. Averaging is a contraction, so
Measured on an Apple M4 Max. Every optimised path is verified to produce output identical to the unoptimised path before it is timed.
| Optimisation | Workload | Before | After | Speedup |
|---|---|---|---|---|
| Operator fusion | 16M elements | 8.1ms | 3.4ms | 2.36x |
| Blocked, threaded matmul | 768 x 768 | 348.0ms | 3.8ms | 91.2x |
| Both, 4-layer MLP | 256 x 256 | 43.0ms | 1.3ms | 33.5x |
The same 2·n³ arithmetic operations throughout. Every gain came from moving memory differently, or from using more cores.
| Stage | Time | GFLOPS | This step | Cumulative |
|---|---|---|---|---|
| Naive triple loop | 349.3ms | 3 | 1.0x | |
| Loop reordering | 28.6ms | 32 | 12.2x | 12.2x |
| Cache blocking | 32.2ms | 28 | 0.89x | 10.8x |
restrict |
32.2ms | 28 | 1.00x | 10.8x |
| Threading and tuning | 3.8ms | 237 | 8.4x | 91.3x |
Loop reordering was the single largest win, and it changed no arithmetic at all. The naive version walks down a column of the right-hand matrix, jumping a full row stride on every step, so each 64-byte cache line fetched yields one useful float out of sixteen. Swapping the loop order so the innermost loop runs contiguously fixes that, and it is worth 12x on its own.
Cache blocking measured slightly slower in isolation, which was not expected. At 768 x 768 a block of the right-hand matrix already fits comfortably in this machine's L2, so tiling adds loop overhead without saving any traffic. It earns its place anyway, because the tiles are what give the threads independent slices of work.
restrict did nothing measurable. It promises the compiler that the three arrays do not overlap, which sometimes unlocks reordering that aliasing would otherwise block. Here it did not, which tells us aliasing was never the constraint.
Threading gave the remaining 8.4x. The choice of which loop to split follows from where threads would collide: splitting the shared-dimension loop is a data race, because different values of that index accumulate into the same output element, which is what the += means. Splitting the output-row loop is safe, because each thread owns its rows outright.
Tile size then falls out of the thread count. A larger tile means better reuse but fewer tiles, and the thread count cannot exceed the number of tiles, so dividing the matrix height by the number of threads balances the two: every thread gets one equal, contiguous slice and none sits idle. That rule reproduces both sizes found by brute-force tuning without being told them, giving 96 at m=768 and 32 at m=256.
One optimisation is missing from the table because it did not work. Hand-written NEON intrinsics for the inner loop measured 48ms against 32ms for the plain scalar version. Clang was already auto-vectorising that loop at width 4 with an interleave of 4, and writing intrinsics constrained its scheduler more than it helped. Unrolling to four independent accumulator chains did not recover the difference either. The scalar loop is kept deliberately.
BLAS is called through NumPy. On this machine that is Apple's Accelerate, which uses hand-tuned assembly and the AMX matrix coprocessor.
| Size | Naive | Forge | BLAS | Forge GFLOPS | BLAS GFLOPS | BLAS faster by |
|---|---|---|---|---|---|---|
| 256 x 256 | 11.0ms | 0.3ms | 0.12ms | 109 | 270 | 2.5x |
| 512 x 512 | 97.1ms | 1.6ms | 0.67ms | 172 | 401 | 2.3x |
| 768 x 768 | 348.0ms | 3.9ms | 2.58ms | 234 | 351 | 1.5x |
234 GFLOPS against 351, in portable C with pthreads, against hand-tuned assembly with dedicated matrix silicon behind it.
Two things stand between the two numbers. A register-blocked micro-kernel, which computes a whole output tile in registers rather than accumulating through memory, is worth perhaps 1.3x and is genuinely available. The AMX matrix units are not: Apple never published how to address them, so no portable C can reach them at all. Part of that 1.5x is not closable from here.
pip install forge-dlPython 3.10 or later. The core library has no dependencies.
Optional extras:
pip install -e ".[plots]" for the option pricing charts, ".[bench]" for the
compiler's BLAS comparison, ".[metal]" for the Apple GPU backend.
from forge import Tensor
a = Tensor([[1.0, 2.0], [3.0, 4.0]])
b = Tensor([[5.0, 6.0], [7.0, 8.0]])
print(a + b) # element-wise addition
print(a @ b) # matrix multiplication
print(a.T) # transpose
x = Tensor([2.0, 3.0], requires_grad=True)
y = ((x * x) + x).sum()
y.backward()
print(x.grad) # dy/dx = 2x + 1from forge.nn import GPT
from forge.nn import CrossEntropyLoss
from forge.optim import Adam
model = GPT(vocab_size=65, embed_dim=128, num_heads=4,
ff_dim=256, num_layers=4, seq_len=32)
criterion = CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=0.002)
logits = model([[1, 2, 3, 4]])
loss = criterion(logits, [2, 3, 4, 5])
optimizer.zero_grad()
loss.backward()
optimizer.step()python3 sidequests/pokemon/pokemon_names.py # invent Pokemon names
python3 sidequests/shakespeare/train_shakespeare.py # data-parallel Shakespeare training
python3 sidequests/shakespeare/resume_shakespeare.py # continue from a checkpointpip install matplotlib
python3 sidequests/options/price_report.pyPrints the convergence table and writes prices.png.
python -m forge.compiler.build
python forge/compiler/benchmark.pyEvery operation records itself in a graph as it runs. Calling backward() walks that graph in reverse topological order, applying the chain rule at each node and accumulating gradients into the leaves.
For a scalar loss
Each Function supplies only the local Jacobian-vector product
Adding a new differentiable operation means subclassing Function and writing forward and backward. Anything composed from existing operations gets its gradient for free, which is why LayerNorm has no backward method of its own.
Every operation is checked against a numerical gradient computed by central finite differences:
from forge.autograd import grad_check
from forge.dtype import float64
def mse(pred):
target = Tensor([1.0, 2.0, 3.0], dtype=float64)
return ((pred - target) ** 2).mean()
pred = Tensor([1.5, 2.5, 3.5], dtype=float64, requires_grad=True)
assert grad_check(mse, [pred])The check compares the analytical gradient against the central difference
Expanding both terms as Taylor series about
This caught a bug where a manual tensor slice in the attention path had silently detached the embedding layer from the graph. The model trained, the loss fell, and the embedding never moved.
Each block is pre-norm with residual connections, so for a block
Writing it as
The identity term keeps that product from collapsing to zero, which is what makes depth trainable.
Attention. With
The scaling by
The mask
Softmax, evaluated in the shift-invariant form
which is algebraically identical to the naive form but never overflows, since every exponent is at most
LayerNorm, over the
with
GELU, in the tanh approximation
approximating the exact
Objective. Training minimises the mean cross-entropy over the
Adam then updates each parameter by
The bias correction matters most early on:
For
costing
Using the wrong stride is invisible when
Matmul picks a path by problem size:
| Condition | Backend |
|---|---|
| float32, above the work threshold | Apple Metal GPU, through MPSMatrixMultiplication |
| float32, below it | Apple Accelerate BLAS, through cblas_sgemm |
| anything else | pure Python triple loop |
The threshold exists because GPU dispatch has a fixed setup cost. For the small matrices in a character-level model, that cost is larger than the multiplication itself, so the GPU is slower than the CPU. The library measures the work and routes accordingly.
flowchart LR
A["Python expression<br/>relu(a + b) * c"] --> B["Graph<br/>nodes, not numbers"]
B --> C["Fusion pass<br/>which ops share a pass"]
C --> D["Code generation<br/>write C, run gcc"]
D --> E["ctypes<br/>load the .so"]
E --> F["Execute<br/>one call per group"]
Five stages, each in its own file under forge/compiler/:
| Stage | File | What it does |
|---|---|---|
| Graph | graph.py |
Records operations as nodes instead of executing them |
| Interpreter | interpreter.py |
Runs the graph one node at a time, as a baseline |
| Fusion | fusion.py |
Groups element-wise chains that can share one memory pass |
| Code generation | codegen.py |
Writes C for each group, compiles it, loads it |
| Execution | compiled_run.py |
Runs the compiled plan |
The fusion rule is that an intermediate can be absorbed into a group only if exactly one operation consumes it. A value read in two places has to exist in memory, so it cannot dissolve into a register.
For relu(a + b) * c the compiler writes this, and nothing else in the project wrote it:
void fused_kernel(const float* in0, const float* in1, const float* in2,
float* out, int n){
for (int i = 0; i < n; i++) {
float t3 = (in0[i] + in1[i]);
float t4 = (t3 > 0.0f ? t3 : 0.0f);
out[i] = (t4 * in2[i]);
}
}Three kernels become one. Three passes over memory become one. The intermediates t3 and t4 live in registers and never reach main memory.
Why fusion helps, precisely. Consider a chain of
Fused into a single loop, every intermediate lives in a register and only the true inputs and the final output touch memory:
The arithmetic is unchanged at
By the roofline model, attainable performance is
Why blocking helps. Cache blocking does not change the
against
Neither transformation improves asymptotic time complexity. Both are constant-factor improvements in the memory hierarchy, and on modern hardware that is where the performance is, because arithmetic is rarely the limit.
Building this from scratch meant every bug was mine to find. Three were instructive enough to write down.
The embedding layer that never trained
Training ran, loss fell, output improved. The embedding weights never moved.
A manual tensor slice in the attention path was copying values into a fresh tensor rather than going through a tracked operation. The autograd graph was severed at that point, so gradients flowing backwards stopped there and never reached the embedding. The rest of the network compensated well enough to hide it.
Nothing about the training curve suggested a problem. It surfaced only when checking whether every layer was actually receiving a gradient, and the embedding's was empty. The fix was making the slice a proper differentiable operation.
The model obsessed with the letter Z
A name generator trained on 1,024 Pokemon names produced almost exclusively names beginning with Z.
The cause was not the model. The training data was sorted alphabetically and never shuffled, so every epoch ended on the z-names and the final gradient steps of each pass pulled the weights toward z openings. Over twenty epochs that bias compounded.
An earlier run on the first 500 names, which end around m, showed the same effect with different letters, which is what identified it. Shuffling the order each epoch fixed it in one line.
The matmul that only worked on square matrices
Flat array indexing means each matrix has its own row stride, and that stride is its own column count. Using the wrong one reads the wrong values, or reads past the end of the array.
Square test data hides this completely, because when every dimension is equal, every stride is the same number. Two separate stride bugs passed all the square tests and produced values around 1e35 on the first non-square input.
The test suite now uses deliberately non-square shapes such as 2x3 @ 3x4, where all three dimensions differ.
The subtraction that ran backwards
Pricing a barrier option needs a flag that is 1 for paths which survived and 0 for those knocked out. The obvious way to write it is 1.0 - breached.
That returned the wrong sign, and the option came out with a negative price.
When Python evaluates 1.0 - tensor, the float does not know how to subtract a tensor, so it hands the job back to the tensor through __rsub__. The implementation computed self - other rather than other - self, silently reversing the operands. Addition and multiplication were unaffected because they commute, so only subtraction and division could expose it.
The lesson is that reflected operators are the one place where writing the obvious implementation gives the wrong answer, precisely because the arguments arrive swapped.
Typed arrays over Python lists. A Python float object costs 24 bytes. A float32 in an array costs 4. For a model with hundreds of thousands of parameters that ratio matters.
Define-by-run for the library, define-then-run for the compiler. The library builds its graph as operations execute, which allows ordinary Python control flow in a forward pass. The compiler needs the opposite, because it has to see the whole computation before anything runs in order to find work worth fusing.
Per-head projections in multi-head attention. The standard implementation makes one large projection and reshapes it into heads. Reshape in this library does not preserve gradients, so each head has its own smaller Query, Key and Value layers instead. Mathematically identical, and every step stays differentiable.
A large negative number instead of negative infinity in the causal mask. Softmax subtracts the row maximum for numerical stability. With true negative infinity that subtraction produces NaN. Using -1e30 gives the same effect without the arithmetic hazard.
The core library is pure Python, so it is orders of magnitude slower than a production framework. That is the point of the exercise, but it is worth stating plainly.
The GPT processes one sequence at a time. Batching was implemented and verified but is not in this branch.
The compiler handles the forward pass over six operations, with no autograd. Only matmul is threaded; the element-wise kernels are single-threaded, which is fine because they are memory-bound rather than compute-bound.
Metal and Accelerate backends require macOS and PyObjC. Everything falls back to pure Python elsewhere.
forge/
tensor.py Tensor class, broadcasting, operator overloading
dtype.py float32 and float64 definitions
serialization.py save and load weights
autograd/
engine.py Function base class, the autograd core
operations.py differentiable operations and their gradients
grad_check.py numerical gradient verification
fusion.py graph-level Linear plus ReLU fusion
mps_backend.py Apple Metal GPU matmul
accelerate_backend.py Apple Accelerate BLAS matmul
nn/
module.py Module base class, parameter registration
layers.py Linear, Embedding, LayerNorm, MultiHeadAttention, GPT
losses.py MSE, BCE, CrossEntropy
parameter.py trainable tensor wrapper
optim/
optimizer.py SGD with momentum, Adam
compiler/
graph.py computation graph representation
fusion.py fusion planning pass
codegen.py C source generation and run-time compilation
interpreter.py baseline node-by-node execution
compiled_run.py compiled plan execution
kernels.c hand-written C kernels, including blocked matmul
build.py builds kernels.so
benchmark.py performance measurement
sidequests/
shakespeare/
train_shakespeare.py data-parallel training across CPU cores
resume_shakespeare.py continue training from a checkpoint
pokemon/
pokemon_names.py character-level name generation
options/
options.py Black-Scholes, Monte Carlo, Asian and barrier pricing
price_report.py convergence table and charts
Built by Srihari Srinivasan, with love :).