TorchLens for teaching how networks work¶

Teaching needs a picture of the whole computation, not a module list: every operation the forward pass actually ran, in order, with the shapes flowing between them, so a student can follow a residual connection, an attention head or a recurrent loop from input to output. TorchLens draws that picture from one call on any PyTorch model, lets a student click into any single operation to see its inputs, outputs and source line, and does the same for a 1943 threshold-neuron network and a modern transformer alike.

The lists below start with drawing and inspecting, then move to the tutorial notebooks and the Model Menagerie. Every link is to a file in this repository or to the Menagerie site.

Start here¶

See the graph and inspect any operation: one call draws every operation the forward ran, including ones no module returns; one lookup gives any operation's output, module and arguments.
Visualize notebook: render a model graph and learn to read the layer labels.
Tutorial notebooks: eleven short notebooks on laptop-sized models, each one rung up from the last. Notebook 01 is "capture and look", 02 reads one activation like a dict, 03 draws the graph when the summary is too long to read. They run on a laptop CPU in about a minute each, or in the browser on Colab, so they work as lab exercises.
TorchLens in 10 minutes: trace, intervene, run a backward pass, see control flow drawn as IF, THEN and ELSE, and compare two traces.

Drawing the graph¶

How TorchLens groups operations into layers: why linear_1 with passes linear_1:1 and linear_1:2 is the same layer run twice, the thing students most often need explained about recurrent and weight-tied models.
Smart-collapse visual reference: how a large graph folds into readable blocks, with rendered examples at each level.
Premade views: trace.draw(view=...) views that each answer one question about the model.
Transformer pictures: token-labelled attention heatmaps and grids and the attention atlas, for the attention lecture.
Notebook cards and the treescope bridge: interactive cards for a notebook walkthrough.
The Netron export and Google Model Explorer export: open the captured graph in a viewer students may already know.

Inspecting every operation¶

Extracting source code with TorchLens: the call stack and the actual source line of every operation, so a drawn node points back at the line of PyTorch that produced it.
Conditional branch attribution: which Python if branch each operation ran under.
Containers: how tuples, dicts and lists flow between operations.
Facets and the facets tutorial notebook: one activation split into its parts, such as an LSTM's output and hidden states or an attention layer's individual heads.
Backward capture and the backward tutorial notebook: draw the autograd graph and connect each backward node to its forward layer, for the backpropagation lecture.
Gradients tutorial notebook: what drives one unit.
Receptive and projective fields: which input positions can affect a given position, and which positions it can affect, computed on the graph that actually ran.

Hands-on exercises¶

Change and rerun: does this layer matter? Zero it and watch the output change.
Intervention notebook: zero-ablate one layer and replay the downstream computation, in a few lines.
Pluck your first activation: one activation out of a model in three lines.
Compare two runs: the same model on two inputs, diffed operation by operation.
Hugging Face tutorial notebook: trace a real pretrained transformer and draw it.

The Model Menagerie¶

The Model Menagerie: a browsable gallery of over 6,600 neural network architectures, from McCulloch and Pitts's 1943 network of threshold neurons to current models, each with its full computational graph drawn by TorchLens. For a course it is a ready-made source of worked drawings: show the same operation vocabulary across eight decades of designs.
The same three lines on any architecture: the trace, draw and edit calls on six architecture families from the Menagerie, the shortest proof that the tool works on whatever a student brings.
Model summaries: a one-call table of every layer with shapes and parameter counts, for the first lecture before the graph.