Lecture 5: Automatic Differentiation Implementation

What this lecture covers

The previous lecture derived reverse-mode automatic differentiation on paper. This lecture switches to an interactive code review of needle ("Necessary Elements of Deep Learning"), the roughly one-thousand-line Python framework students extend in the course homeworks. Rather than presenting new theory, it walks through the actual data structures that make a computational graph runnable: how a tensor stores its history, how operations attach themselves to that history, and how gradients are computed one step at a time.

By the end, you can explain how a Tensor object relates to the underlying Value and Op classes, trace a computational graph through its inputs and op fields, and understand the mechanics your homework's reverse-mode AD implementation will build on.

Key ideas

• Value class: the base class behind Tensor; it stores cached_data (the array result), inputs (the values feeding an operation), op (which operation produced it), and requires_grad.
• Op class: each operation (such as EWiseAdd or AddScalar) is a subclass of Op that implements a compute method (runs the array math) and a gradient method (produces adjoints for its inputs).
• Computational graph as data: a tensor is not just an array; its inputs and op fields let you trace back through every operation that produced it, forming a directed acyclic graph.
• Lazy vs eager evaluation: eager mode (needle's default, like PyTorch) computes cached_data immediately when an operation runs; lazy mode defers computation until the data is actually needed, which can help when building large graphs like PyTorch/XLA does on TPUs.
• detach: strips a tensor's computational graph history, returning a plain value; used to avoid accidentally accumulating an entire chain of operations (and the memory that goes with it) across training iterations.
• realize_cached_data: the function that actually triggers computation, recursively realizing each input's cached data before calling op.compute.
• Gradient functions operate on tensors, not arrays: a gradient function takes the output adjoint and the forward node (both Tensor objects) so that the resulting adjoints are themselves part of a computational graph, enabling gradient-of-gradient computations.

Walkthrough

Setting up and exploring needle in Colab (0:00)

The lecture opens with environment setup: cloning the course repo into Google Drive, symlinking it into a Colab content folder, and adding the needle package to the Python path. The needle library is introduced as roughly three files — __init__.py, autograd.py (about 400 lines defining the core data structures) and ops.py (about 300 lines of operator definitions). A few basic operations are demonstrated: creating a tensor with ndl.Tensor, adding a scalar to it, and reading back .shape, .dtype, and .numpy().

The Value and Tensor data structures (11:15)

Tensor is shown to be a thin subclass of Value, which carries the real fields: cached_data (an NDArray, currently backed by numpy), inputs (a list of the values feeding this computation), op (the operation that produced it), and requires_grad. Using the exponential-plus-one graph from the prior lecture as a running example, the lecture shows how v4.inputs traces back to v2 and v3, and ultimately to the leaf node v1, whose op is None.

The Op class and building a graph by hand (16:22)

Ops such as EWiseAdd and AddScalar are subclasses of TensorOp. Each defines compute (the array-level math) and gradient (the adjoint calculation). The lecture constructs v1 through v4 directly in code, confirming that v4.inputs, v4.op, and v4.cached_data match the expected computational graph, and writes a small print_node helper that prints a node's Python id, its inputs' ids, its operation type, and its data — a convenient way to inspect a graph.

What happens when you write x1 + x2 (29:53)

Tracing operator overloading, the lecture follows x1 + x2 through Tensor.__add__, into EWiseAdd.__call__ (inherited from TensorOp), and into Tensor.make_from_op, which uses Tensor.__new__ plus an internal _init method to attach op and inputs without yet computing anything. The actual arithmetic only happens inside realize_cached_data, which recursively realizes each input and then calls op.compute.

Lazy mode versus eager mode (37:17)

Toggling needle.autograd.LAZY_MODE shows that in lazy mode cached_data stays None until something (printing, .numpy(), .data) forces evaluation via detach and realize_cached_data. The lecture contrasts eager execution (PyTorch's default, where graph construction and computation are interleaved) with lazy execution (useful when graph-construction overhead is negligible next to batched computation, as in PyTorch XLA on TPUs).

The memory trap of accumulating a computational graph (44:26)

A hundred-iteration loop that accumulates sum_loss += x * x is used to show that each iteration silently extends a long chain of computational-graph nodes rather than a single scalar, which is a common source of runaway memory use in training loops. Calling .detach() (or PyTorch's no-grad context) strips the graph, leaving a plain cached value with empty inputs and no op.

Gradient functions and reverse-mode AD mechanics (53:50)

Returning to v4 = v2 * v3, the lecture derives the adjoint rules v_(2->4) = v4's adjoint * v3 and v_(3->4) = v4's adjoint * v2 by hand, then matches them against EWiseMul's gradient implementation. Plugging in a Tensor(1.0) output adjoint and running the gradient function on v4 reproduces these values numerically, confirming that each op's gradient method defines a single reverse-mode step. Because gradients are themselves tensors carrying their own computational graph, the resulting adjoints (such as v3adjoint) can be differentiated again, enabling gradient-of-gradient computations.

Before you watch

• Watch the previous lecture on deriving numerical and reverse-mode automatic differentiation, since this lecture assumes that derivation and reuses its computational-graph example.
• Be comfortable reading Python classes, __init__/__new__, and basic numpy array operations, since the lecture reads needle's source code directly.

Check your understanding

1. What information does a Value object store besides its computed array, and why does it need that information?
2. Walk through what happens, step by step, when you write x1 + x2 on two needle tensors, from operator overloading to realize_cached_data.
3. Why can accumulating a loss inside a training loop without calling detach lead to excessive memory use?
4. In the EWiseMul gradient function for v4 = v2 * v3, what are the adjoints for v2 and v3, and why does each depend on the other input?
5. Why do gradient functions operate on Tensor objects rather than raw arrays, and what capability does that give the framework?