How to Implement Autograd from Scratch in Python: A Complete Guide to TinyTorch's Automatic Differentiation Engine

You can implement autograd from scratch in Python by defining a Function base class that records operation inputs, creating concrete backward classes (like AddBackward and MatmulBackward) that compute local gradients using the chain rule, and monkey-patching tensor operations to build a dynamic computation graph that supports reverse-mode automatic differentiation.

When building deep learning frameworks, automatic differentiation (autograd) is the engine that powers gradient-based optimization. If you want to implement autograd from scratch in Python, the TinyTorch library in the harvard-edge/cs249r_book repository provides a complete educational implementation. This article walks through the architecture, showing exactly how reverse-mode automatic differentiation works without relying on external autograd libraries.

The Core Architecture: Function Base Class and Backward Operations

The foundation of TinyTorch's autograd system is the Function abstraction defined in tinytorch/src/06_autograd/06_autograd.py. This base class establishes the contract for all differentiable operations.

class Function:
    """
    Base class for differentiable operations.
    Each subclass stores the tensors it needs (saved_tensors) and implements
    `apply(grad_output)` to return the gradient for each input.
    """
    def __init__(self, *tensors):
        self.saved_tensors = tensors

    def apply(self, grad_output):
        raise NotImplementedError

Location: 06_autograd.py – Function definition

Each operation that supports backpropagation inherits from Function and implements the apply() method to compute local gradients using the chain rule.

Concrete Backward Classes: Implementing the Chain Rule

TinyTorch implements specific backward classes for every arithmetic and neural network operation. Each class overrides apply() to return the gradient with respect to each input tensor.

Operation	Backward class	Core logic
Addition	AddBackward	Returns grad_output for each input (handles broadcasting).
Multiplication	MulBackward	Returns grad_output * other_input.
Matrix multiplication	MatmulBackward	Uses grad_output @ B.T and A.T @ grad_output.
ReLU	ReLUBackward	Multiplies grad_output by the mask (x > 0).
Sigmoid	SigmoidBackward	Uses saved sigmoid output σ(x)*(1-σ(x)).
Softmax	SoftmaxBackward	Implements softmax * (grad_output - Σ(grad_output*softmax)).
Cross‑entropy loss	CrossEntropyBackward	Computes (softmax - one_hot) / batch_size.

Location examples:

• AddBackward implementation
• MatmulBackward implementation
• CrossEntropyBackward implementation

Enabling Autograd: Monkey-Patching Tensor Operations

The enable_autograd() function in tinytorch/src/06_autograd/06_autograd.py activates the autograd system without modifying the original Tensor class source code. This approach uses monkey-patching to wrap arithmetic methods.

def enable_autograd(quiet=False):
    # 1️⃣ Extend Tensor.__init__ to create .requires_grad and .grad attributes

    # 2️⃣ Save original arithmetic methods (add, mul, matmul, etc.)

    # 3️⃣ Replace them with "tracked" versions that:

    #    • Call the original operation to obtain a result Tensor.

    #    • If any operand has requires_grad=True, wrap the result in a new

    #      Function subclass, store the backward graph, and set

    #      result.requires_grad = True.

    # 4️⃣ Add a Tensor.backward() method that:

    #    • Starts from the leaf Tensor, seeds grad_output = 1 (or given)

    #    • Walks the saved graph in reverse, invoking each Function.apply()

    #    • Accumulates gradients into .grad of each Tensor.

Location: enable_autograd definition

Key implementation details:

• Monkey-patching preserves the original educational Tensor class for earlier labs while making autograd optional.
• Gradient-aware initialization adds .requires_grad and .grad attributes to every Tensor created after activation.
• Safety helpers (_get_requires_grad, _ensure_grad_attrs) ensure compatibility with tensors instantiated before enable_autograd() is called.

The Backward Pass: Traversing the Computation Graph

The backward() method added to Tensor performs reverse-mode automatic differentiation by traversing the computation graph stored during the forward pass.

def backward(self, grad_output=None):
    if grad_output is None:
        grad_output = np.ones_like(self.data)  # seed for scalar loss

    self.grad = grad_output
    # Walk the graph stored in self._grad_fn (the Function that produced this Tensor)

    # Recursively call fn.apply() → grads for parents → parent.backward(grads)

Location: Inside tinytorch/src/06_autograd/06_autograd.py, implemented within the enable_autograd() setup (see lines around the tracked operations section).

The graph is stored lazily: each Tensor created by a tracked operation receives an attribute _grad_fn pointing to the Function instance that knows its parent tensors.

Practical Usage Examples

Example 1 – Simple Scalar Arithmetic

from tinytorch.core.autograd import enable_autograd
from tinytorch.core.tensor import Tensor

enable_autograd()                 # Activate the autograd engine

x = Tensor(2.0, requires_grad=True)
y = Tensor(3.0, requires_grad=True)

z = x * y + x                     # (2 * 3) + 2 = 8

z.backward()                      # dz/dx = y + 1 = 4, dz/dy = x = 2

print("dz/dx =", x.grad)         # → [4.0]

print("dz/dy =", y.grad)         # → [2.0]

Source reference: The operations (__mul__, __add__) are the tracked versions defined in enable_autograd() (see lines ~200‑240 of 06_autograd.py).

Example 2 – Linear Layer with ReLU

import numpy as np
from tinytorch.core.autograd import enable_autograd
from tinytorch.core.tensor import Tensor
from tinytorch.core.nn import Linear, ReLU   # thin wrappers that use Tensor ops

enable_autograd()

# Random weight matrix (2×3) and bias (2,)

W = Tensor(np.random.randn(2, 3), requires_grad=True)
b = Tensor(np.random.randn(2), requires_grad=True)

linear = Linear(W, b)
relu = ReLU()

x = Tensor(np.random.randn(3), requires_grad=True)   # input vector

out = relu(linear(x))                               # forward pass

loss = out.sum()                                    # scalar loss

loss.backward()

print("grad w.r.t. W:", W.grad.shape)   # (2,3)

print("grad w.r.t. b:", b.grad.shape)   # (2,)

print("grad w.r.t. x:", x.grad.shape)   # (3,)

Key files:

• Linear layer implementation: tinytorch/src/07_optimizers/07_optimizers.py (uses Tensor.matmul).
• ReLU activation: tinytorch/src/09_convolutions/09_convolutions.py (calls ReLUBackward).

Example 3 – Cross-Entropy Loss for Classification

from tinytorch.core.autograd import enable_autograd
from tinytorch.core.tensor import Tensor
from tinytorch.core.nn import Linear
from tinytorch.core.losses import CrossEntropy   # uses CrossEntropyBackward

enable_autograd()

logits = Linear(Tensor(np.random.randn(4, 10), requires_grad=True),
                Tensor(np.zeros(10), requires_grad=True))(Tensor(np.random.randn(4, 10)))
targets = Tensor([2, 0, 1, 3])      # class indices

loss = CrossEntropy(logits, targets)
loss.backward()

print("logits grad shape:", logits.grad.shape)  # (4,10)

Source reference: Cross-entropy backward logic is in CrossEntropyBackward (see lines 48‑66 of 06_autograd.py).

Key Files in the Repository

File	Role	GitHub Link
tinytorch/src/06_autograd/06_autograd.py	Core autograd engine: Function, all backward classes, enable_autograd(), and Tensor.backward()	https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/src/06_autograd/06_autograd.py
tinytorch/__init__.py	Package entry-point; imports and invokes enable_autograd() automatically	https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/__init__.py
tinytorch/core/tensor.py	Original Tensor implementation (operations, data storage)	https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/core/tensor.py
tinytorch/tests/06_autograd/	Unit-tests for each backward class and end-to-end gradient flow	https://github.com/harvard-edge/cs249r_book/tree/dev/tinytorch/tests/06_autograd
tinytorch/src/07_optimizers/07_optimizers.py	Example usage of autograd in optimizers (SGD, Adam)	https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/src/07_optimizers/07_optimizers.py
tinytorch/src/09_convolutions/09_convolutions.py	Demonstrates autograd with convolutional layers and ReLU	https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/src/09_convolutions/09_convolutions.py
tinytorch/src/11_embeddings/11_embeddings.py	Shows autograd with embedding look-ups	https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/src/11_embeddings/11_embeddings.py

These files together illustrate a complete, from-scratch automatic differentiation pipeline that can be studied, modified, or extended for educational or research purposes.

Summary

To implement autograd from scratch in Python, you need four essential components working together:

• A Function base class that stores input tensors (saved_tensors) and defines an apply() interface for computing local gradients.
• Concrete backward classes (e.g., AddBackward, MatmulBackward, CrossEntropyBackward) that implement the chain rule for specific operations.
• Monkey-patching via enable_autograd() to intercept tensor operations and build a dynamic computation graph without modifying the original Tensor source.
• A backward() traversal method that walks the computation graph in reverse, invoking each stored Function.apply() and accumulating gradients into tensor.grad.

The TinyTorch implementation in tinytorch/src/06_autograd/06_autograd.py demonstrates that a fully functional autograd engine requires only a few hundred lines of Python while remaining powerful enough to train neural networks with optimizers, convolutions, and embeddings.

Frequently Asked Questions

What is the Function base class in TinyTorch autograd?

The Function class is the abstract base that enables automatic differentiation in TinyTorch. It stores input tensors in saved_tensors during the forward pass and requires subclasses to implement apply(grad_output), which computes the gradient of the loss with respect to each input using the chain rule. This design pattern separates the forward computation from the backward logic, making it easy to add new differentiable operations.

How does enable_autograd() work without modifying the original Tensor class?

The enable_autograd() function uses monkey-patching to add autograd capabilities to the existing Tensor class dynamically. It saves references to the original methods like __add__ and __mul__, then replaces them with wrapped versions that create Function nodes and set requires_grad flags when needed. This approach keeps the core Tensor implementation clean for educational purposes while allowing optional activation of the full autograd engine.

How does the backward pass traverse the computation graph?

The Tensor.backward() method initiates reverse-mode differentiation by seeding the gradient with ones (or a provided grad_output) and then walking the graph stored in _grad_fn attributes. It recursively calls Function.apply() for each operation in reverse topological order, passing gradients backward through the chain rule and accumulating results into each tensor's .grad attribute. This traversal efficiently computes gradients for all parameters in a single pass from the loss node to the inputs.

Can TinyTorch autograd handle higher-order derivatives?

The current TinyTorch implementation focuses on first-order automatic differentiation for educational clarity. While the architecture could theoretically support higher-order derivatives by making the backward() operation itself differentiable (requiring Function objects that track their own gradient computations), the existing code in tinytorch/src/06_autograd/06_autograd.py is optimized for standard neural network training with scalar losses. For higher-order gradients, you would need to extend the Function classes to return Tensor objects with their own grad_fn references during the backward pass.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}