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

> Implement autograd from scratch in Python using a Function base class and monkey-patching tensor operations. Learn how to build a dynamic computation graph for reverse-mode automatic differentiation.

- Repository: [Harvard Edge Computing/cs249r_book](https://github.com/harvard-edge/cs249r_book)
- Tags: how-to-guide
- Published: 2026-02-19

---

**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`](https://github.com/harvard-edge/cs249r_book/blob/main/tinytorch/src/06_autograd/06_autograd.py). This base class establishes the contract for all differentiable operations.

```python
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](https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/src/06_autograd/06_autograd.py#L33-L45)

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](https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/src/06_autograd/06_autograd.py#L33-L49)
- [MatmulBackward implementation](https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/src/06_autograd/06_autograd.py#L63-L71)
- [CrossEntropyBackward implementation](https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/src/06_autograd/06_autograd.py#L48-L63)

## Enabling Autograd: Monkey-Patching Tensor Operations

The `enable_autograd()` function in [`tinytorch/src/06_autograd/06_autograd.py`](https://github.com/harvard-edge/cs249r_book/blob/main/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.

```python
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](https://github.com/harvard-edge/cs249r_book/blob/dev/tinytorch/src/06_autograd/06_autograd.py#L96-L140)

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.

```python
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`](https://github.com/harvard-edge/cs249r_book/blob/main/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

```python
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`](https://github.com/harvard-edge/cs249r_book/blob/main/06_autograd.py)).

### Example 2 – Linear Layer with ReLU

```python
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`](https://github.com/harvard-edge/cs249r_book/blob/main/tinytorch/src/07_optimizers/07_optimizers.py) (uses `Tensor.matmul`).
- ReLU activation: [`tinytorch/src/09_convolutions/09_convolutions.py`](https://github.com/harvard-edge/cs249r_book/blob/main/tinytorch/src/09_convolutions/09_convolutions.py) (calls `ReLUBackward`).

### Example 3 – Cross-Entropy Loss for Classification

```python
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`](https://github.com/harvard-edge/cs249r_book/blob/main/06_autograd.py)).

## Key Files in the Repository

| File | Role | GitHub Link |
|------|------|-------------|
| [`tinytorch/src/06_autograd/06_autograd.py`](https://github.com/harvard-edge/cs249r_book/blob/main/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`](https://github.com/harvard-edge/cs249r_book/blob/main/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`](https://github.com/harvard-edge/cs249r_book/blob/main/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`](https://github.com/harvard-edge/cs249r_book/blob/main/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`](https://github.com/harvard-edge/cs249r_book/blob/main/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`](https://github.com/harvard-edge/cs249r_book/blob/main/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`](https://github.com/harvard-edge/cs249r_book/blob/main/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`](https://github.com/harvard-edge/cs249r_book/blob/main/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.