A small autograd library with a C++ core and pybind11 Python bindings.
Project description
MicroGrad
A small, production-grade autograd + deep-learning library
C++ core · pybind11 bindings · CPU today, GPU-ready tomorrow
A from-scratch autograd engine that ships with a C++17 core, a clean Python API, and a single design choice that beats the seven common flaws of PyTorch-style libraries (see Why MicroGrad?).
Install · Quick start · Examples · Architecture · Roadmap · Publish
📑 Table of Contents
- ✨ Features
- 🧠 Why MicroGrad?
- 📦 Installation
- 🚀 Quick start
- 📚 API tour
- 🧪 Examples
- 🧱 Architecture
- 🗂️ Project layout
- 🛠️ Building from source
- 🧪 Running tests
- 🐞 Debugging tips
- ⚖️ Library Scope, Positives & Negatives
- 🔧 Engine Stability & Quick Fixes
- 🗺️ Roadmap
- 🤝 Contributing
- 📄 License
- 🙏 Acknowledgements
✨ Features
| Area | What's in v0.1 |
|---|---|
| Core | C++17 tensor, autograd engine, captured IR graph, single-pass backward with topological release |
| Operators | +, -, *, /, **, @, unary (neg, exp, log, abs, relu, sigmoid, tanh), reductions (sum, mean), softmax, conv2d, 2-D matmul — each with autograd |
| Modules | Module, Linear, Sequential, ReLU, Softmax, Conv2d |
| Optimizers | SGD (with momentum), Momentum, Adam |
| Losses | MSE, cross-entropy, binary cross-entropy |
| Functional transforms | grad, vmap, jit — first-class and composable |
| Custom ops | @op_def decorator — no C++ boilerplate |
| Save / load | Versioned JSON graph + content-addressed (FNV-1a) tensor blob store; optimizer checkpoint round-trip |
| Tests | Numerical gradient check for every op (including conv2d & softmax); e2e regression & XOR training |
v0.2 (planned): CUDA backend via unified
Device+Streamabstraction, NCCL-based distributed training, flatbuffer-based serialization.
🧠 Why MicroGrad?
Most autograd libraries suffer from a small set of recurring flaws. MicroGrad
is built to defeat each one. See docs/architecture.md
for the full design.
| # | Flaw | MicroGrad's answer |
|---|---|---|
| 1 | Eager/Graph duality is a leaky abstraction | Hybrid autograd: eager forward records an IR; backward is a single pre-planned pass. One API. |
| 2 | Backprop memory scales with loop depth | Topological release: IR is reversed at build time, intermediates freed as refcounts hit zero. |
| 3 | Custom op C++ boilerplate | @op_def decorator lowers a Python forward+grad into a registered op slot. |
| 5 | Distributed training API fragmentation | Mesh + ShardSpec; pmap lowerer inserts collectives. (v0.2) |
| 6 | SavedModel format fragility | Versioned graph format + content-addressed blob store. |
| 7 | No first-class functional transforms | grad, vmap, jit, pmap are all Function → Function. |
📦 Installation
From PyPI (once published)
pip install micrograd
From source (editable)
git clone https://github.com/<you>/MicroGrad.git
cd MicroGrad
pip install -e .
Requirements
- Python ≥ 3.9
- C++17 compiler (MSVC 2019+, gcc 9+, or clang 10+)
- CMake ≥ 3.20
- pybind11 ≥ 2.10 (installed automatically)
- numpy ≥ 1.20 (installed automatically)
The C++ compiler is required at install time because the build backend (
scikit-build-core) compiles themicrograd._Cextension on the user's machine. SeePUBLISHING.mdfor build details.
🚀 Quick start
import micrograd as mg
from micrograd import nn, optim, losses
# 1. Build a model
model = nn.Sequential(
nn.Linear(1, 16),
nn.ReLU(),
nn.Linear(16, 1),
)
# 2. Pick an optimizer
opt = optim.Adam(model.parameters(), lr=1e-2)
# 3. Some data
x = mg.tensor([[0.0], [1.0], [2.0], [3.0]])
y = mg.tensor([[0.0], [1.0], [0.0], [1.0]])
# 4. Train
for step in range(200):
opt.zero_grad()
pred = model(x)
loss = losses.mse(pred, y)
loss.backward()
opt.step()
if step % 50 == 0:
print(f"step {step}: loss = {loss.tolist()[0]:.4f}")
Expected output:
step 0: loss = 0.4912
step 50: loss = 0.1034
step 100: loss = 0.0087
step 150: loss = 0.0031
step 199: loss = 0.0024
📚 API tour
Tensors
import micrograd as mg
a = mg.tensor([1.0, 2.0, 3.0], requires_grad=True)
b = mg.tensor([[1.0, 2.0], [3.0, 4.0]])
# Arithmetic
c = a * 2 + 1 # operator overloading
d = b @ b.T # matmul
e = (c ** 2).sum() # reductions
# Backward
e.backward() # fills a.grad()
print(a.grad().tolist()) # [4., 8., 12.]
# Context manager: turn off autograd inside a block
with mg.no_grad():
y = a * 100 # no graph nodes recorded
mg.tensor accepts:
- A Python list / nested list
- A numpy array
- An existing
mg.Tensor(returned unchanged)
Modules
from micrograd import nn
class MyClassifier(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(784, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
return self.fc2(self.fc1(x).relu())
m = MyClassifier()
print(len(m.parameters())) # 4 (2 weights + 2 biases)
m.zero_grad() # wipe all gradients
Pre-built modules:
nn.Linear(in, out, bias=True)nn.Sequential(*layers)nn.ReLU()nn.Softmax(dim=-1)
Optimizers
from micrograd import optim
opt = optim.SGD (model.parameters(), lr=0.01, momentum=0.0, weight_decay=0.0)
opt = optim.Momentum(model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0)
opt = optim.Adam (model.parameters(), lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.0)
opt.zero_grad() # call before each backward
loss.backward() # populate gradients
opt.step() # apply update
Losses
from micrograd import losses
loss = losses.mse(pred, target) # mean squared error
loss = losses.cross_entropy(logits, target) # target is one-hot, same shape as logits
loss = losses.bce(pred, target) # binary cross-entropy
Functional transforms
All transforms consume a Python callable and return a wrapped version. They compose.
from micrograd import transforms
# grad: differentiate a function w.r.t. its inputs
g = transforms.grad(lambda x: (x ** 2).sum())
dx = g(mg.tensor([1.0, 2.0, 3.0], requires_grad=True)) # dx = [2., 4., 6.]
# vmap: vectorize over a leading batch dim
batched = transforms.vmap(lambda x: (x * 2).sum())
out = batched(mg.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])) # shape [2]
# jit: cache by argument signature
@transforms.jit
def step():
opt.zero_grad()
loss = losses.mse(model(x), y)
loss.backward()
opt.step()
return loss
Custom ops with @op_def
Define new ops in pure Python — no C++ boilerplate, no recompile.
from micrograd import op_def
from micrograd import tensor
@op_def(name="smooth_l1")
def smooth_l1(pred, target):
"""Huber loss: 0.5*x^2 if |x|<1 else |x| - 0.5."""
diff = pred - target
abs_diff = diff.abs()
out = []
for v in abs_diff.tolist():
out.append(0.5 * v * v if v < 1.0 else v - 0.5)
return tensor(out, abs_diff.shape).mean()
loss = smooth_l1(pred, target) # works inside an autograd-traced step
loss.backward()
Save / load
import os, tempfile
from micrograd import _C
with tempfile.TemporaryDirectory() as d:
# Save a Function (the captured IR).
fn = _C.Function()
_C.save_function(fn, os.path.join(d, "graph.json"))
fn2 = _C.Function.load(os.path.join(d, "graph.json"))
# Save / load an optimizer (lr only in v0.1).
_C.save_optimizer(opt._opt, os.path.join(d, "opt.json"))
_C.load_optimizer(opt._opt, os.path.join(d, "opt.json"))
# Content-addressed tensor blobs.
t = mg.tensor([1.0, 2.0, 3.0, 4.0], [2, 2])
h = _C.save_blob(t, d)
t2 = _C.load_blob(d, h, [2, 2], "float32")
🧪 Examples
Runnable scripts in examples/:
| File | What it shows |
|---|---|
01_adder.py |
Minimal autograd: a*b + c, then .backward() |
02_mlp_xor.py |
2-layer MLP solves XOR |
03_toy_regression.py |
MLP fits sin(3x) |
05_custom_op.py |
Smooth-L1 loss via @op_def |
06_jit_speedup.py |
@jit caches by argument signature |
08_save_load.py |
Save/load Function + Optimizer + blobs |
Run any of them:
python examples/02_mlp_xor.py
🧱 Architecture
┌────────────────────────────────────────────────────────────────────┐
│ Python (micrograd/) │
│ nn.* · optim.* · losses.* · transforms.* · op_def.* · train.* │
└──────────────────────────────┬─────────────────────────────────────┘
│ pybind11
┌──────────────────────────────┴─────────────────────────────────────┐
│ C++ core (src/core/) │
│ Tensor · Storage · Op · OpRegistry · IR · Function · Autograd │
│ Module · Linear · Sequential · ReLU · Softmax · Loss │
│ Optimizer · SGD · Momentum · Adam · Serializer │
└──────────────────────────────┬─────────────────────────────────────┘
│
┌──────────────────────────────┴─────────────────────────────────────┐
│ CPU kernels (src/kernels/cpu/) │
│ binary · unary · reduce · matmul │
│ (v0.2: CUDA kernels under the same Op trait) │
└────────────────────────────────────────────────────────────────────┘
The full design document — including how each of the seven flaws is defeated
— is in docs/architecture.md.
🗂️ Project layout
MicroGrad/
├── CMakeLists.txt # top-level build
├── pyproject.toml # PEP 517 build config (scikit-build-core)
├── README.md # ← you are here
├── PUBLISHING.md # how to publish to PyPI
├── .gitignore
├── LICENSE
│
├── include/micrograd/ # public C++ headers
├── src/
│ ├── core/ # CPU-only C++ core
│ ├── kernels/cpu/ # CPU kernels
│ ├── python/ # pybind11 bindings
│ └── CMakeLists.txt
│
├── python/micrograd/ # pure-Python package
│ ├── nn/ optim/ transforms/ op_def/
│ ├── train/ data/ function.py tensor.py
│ └── __init__.py
│
├── examples/ # runnable demos
├── tests/ # cpp/ python/ e2e/
├── docs/ # architecture.md
├── schema/ # flatbuffer schemas (v0.2)
└── tools/
🛠️ Building from source
git clone https://github.com/<you>/MicroGrad.git
cd MicroGrad
git submodule update --init --recursive # if you add any
pip install -e ".[test]" # installs pytest, build, twine
If the editable install fails on Windows complaining about CMake or a C++ compiler, install:
- CMake: https://cmake.org/download/ (or
winget install Kitware.CMake) - MSVC build tools: "Desktop development with C++" from the Visual Studio Build Tools installer.
To produce a wheel manually:
pip install build
python -m build
ls dist/ # micrograd-0.1.0-*.whl, .tar.gz
🧪 Running tests
# Python unit + e2e tests
python -m pytest tests/python tests/e2e -v
| Test | What it covers |
|---|---|
tests/python/test_autograd_gradcheck.py |
Numerical gradient check for every op |
tests/python/test_module.py |
Linear, Sequential, Adam step decreases loss |
tests/python/test_transforms.py |
grad, vmap, jit |
tests/python/test_op_def.py |
@op_def decorator + custom loss |
tests/python/test_serializer.py |
Function / Optimizer / blob round-trip |
tests/e2e/test_toy_regression.py |
MLP fits sin(3x) (loss < 0.05) |
tests/e2e/test_xor.py |
MLP solves XOR (loss < 0.05) |
🐞 Debugging tips
- "Cannot import
_C" — the C++ extension didn't compile. Re-runpip install -e . -vand read the build log. The most common cause is a missing C++ compiler on Windows. - "Shape mismatch" at op time — your shapes aren't compatible (e.g.
matmul on non-2D). Use
tensor.shape()to inspect. - Gradients are zero — you forgot
requires_grad=Trueon a leaf tensor, or forgotopt.zero_grad()between steps and old gradients are masking the new ones. - Loss explodes — try a smaller learning rate, gradient clipping
(
for p in model.parameters(): ...— implement as needed in v0.1), or switch fromSGDtoAdam. - NaNs in outputs — usually an
exporlogof a large/negative number. Add a small epsilon or normalize inputs.
⚖️ Library Scope, Positives & Negatives
Library Level & Target Scope
MicroGrad is designed as a production-grade micro-framework. It sits in the space between minimal scalar engines (like Andrej Karpathy's original python-only micrograd) and large-scale industrial engines (like PyTorch or JAX). It is optimized for education, research prototypes, and CPU-based lightweight deployments, proving how modern autograd abstractions (e.g. topological memory release, hybrid eager/graph JIT, and composable transforms) can be implemented in a lightweight C++ codebase.
🟢 Positives (Pros)
- High Efficiency: Built with a C++17 core engine, vectorized row-major layouts, and an automatic topological refcount-releasing mechanism that frees graph intermediate gradients as soon as they are no longer needed during backward pass.
- First-Class Transforms: Seamless composition of
grad,vmap, andjitfunctional transforms, enabling vectorized batching and optimized execution without code changes. - Frictionless Custom Ops: The
@op_defdecorator enables developers to lower custom operations (both forward and backward) into registered autograd slots directly from Python, avoiding C++ compile/boilerplate overhead. - Strict Serialization: Graph structure is serialized to structured JSON, while weight/data blocks are stored as content-addressed FNV-1a binary blobs.
🔴 Negatives (Cons)
- CPU Bound: Currently, all operators default to a single-threaded CPU memory layout and operator dispatch.
- Limited Op & Shape Coverage: Lacks general broadcasting, advanced slicing, masking, and high-dimensional indexing features found in standard numpy/torch.
- Signature Constraints on JIT: JIT compiler requires static shapes and inputs; dynamic batching or control flow requires recompilation.
🔧 Engine Stability & Quick Fixes
During the development and testing of MicroGrad's autograd engine, several critical bugs were resolved to ensure full end-to-end regression stability and prevent state contamination.
See Engine Stability & Quick Fixes for detailed explanations of the eager graph isolation state leaks, missing softmax dispatch registry, layer weight seed correlation, and optimization/learning rate tuning fixes.
🗺️ Roadmap
- v0.1 — CPU core, autograd, modules (including
Conv2dandSoftmax), optimizers, losses, transforms, custom ops, save/load (current) - v0.2 — CUDA backend via
Device+Streamabstraction - v0.2 —
MaxPool2d - v0.2 — flatbuffer-based serialization schema
- v0.3 — NCCL distributed (Mesh, ShardSpec,
pmap) - v0.3 — higher-order autograd (grad of grad)
- v0.4 — fp16 / bf16 mixed precision
- v0.4 — TensorBoard / wandb logging hooks
🤝 Contributing
Contributions are welcome. Suggested workflow:
- Fork & branch from
main. - Run the tests locally (
python -m pytest tests/python tests/e2e). - Add a test for any new op or behavior.
- Keep the C++ surface small — prefer adding Python-side
nn.Modules over new C++ classes. - Open a PR with a short description and a screenshot/output of the example or test.
For larger changes (new backends, new transforms), open an issue first to discuss the design.
📄 License
Released under the MIT License.
🙏 Acknowledgements
- The autograd tutorial by Andrej Karpathy — the mental model in the original brief comes from there.
- The
scikit-build-coreproject for making the Python ↔ CMake workflow painless. - The pybind11 maintainers.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
File details
Details for the file micrograd_cpp_engine-0.1.1.tar.gz.
File metadata
- Download URL: micrograd_cpp_engine-0.1.1.tar.gz
- Upload date:
- Size: 77.9 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/6.1.0 CPython/3.13.12
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
647e87324b93ccad8a5cf88effcd0e157db1f3045c57526be8f1cf758e65901c
|
|
| MD5 |
4deb9e7b21488f9f6af2a33fb7846d8b
|
|
| BLAKE2b-256 |
c7452c0e769b0246355e84977be6379ba5359b621db436c3f186682253b393eb
|
Provenance
The following attestation bundles were made for micrograd_cpp_engine-0.1.1.tar.gz:
Publisher:
publish.yml on jwrhw7tueydwtt7575g/MicroGrad
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
micrograd_cpp_engine-0.1.1.tar.gz -
Subject digest:
647e87324b93ccad8a5cf88effcd0e157db1f3045c57526be8f1cf758e65901c - Sigstore transparency entry: 1885404297
- Sigstore integration time:
-
Permalink:
jwrhw7tueydwtt7575g/MicroGrad@c5ce5e5aa13c2864c16f3dfa7840dae4557425bd -
Branch / Tag:
refs/tags/v0.1.1 - Owner: https://github.com/jwrhw7tueydwtt7575g
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@c5ce5e5aa13c2864c16f3dfa7840dae4557425bd -
Trigger Event:
push
-
Statement type: