Inspect, visualize and intervene on any operation in any neural network.
Understanding and controlling the computations inside neural networks is the most important scientific challenge of our time (at least). Progress on this problem, known as mechanistic interpretability, would help make AI more beneficial for society, illuminate how our brains work, and guard against undesirable behaviors like deception or killing us all.
Some aspects of mechanistic interpretability are deep and open-ended. We need to understand dauntingly complex systems that may soon outsmart us, with new frontier models arriving faster than we can study the current ones. It's not even clear what a "solution" would look like, or how many obstacles remain. We'll surely need help from the AI agents themselves.
Other aspects are just plumbing, and while tedious, they are finite: the same small set of tasks for every model, like extracting and editing its activations and gradients. For instance, PyTorch has long lacked a built-in way to extract every intermediate activation from arbitrary networks (an official docstring admits this can take "some massaging"). But all networks are built from the same basic parts, so a few well-made tools should work on all of them.
The modest goal of TorchLens is to fix these plumbing problems once and for all, so you and your agents can spend your time and tokens on the deep questions, and maybe even solve some of them.
Almost everything you do in TorchLens starts with one function, tl.trace. You don't have to set anything up. Pass it any PyTorch model and its input exactly as they are, with no hooks, wrappers or changes to your code:
import torchlens as tl
from transformers import GPT2LMHeadModel
gpt2 = GPT2LMHeadModel.from_pretrained("gpt2").eval()
text = "The capital of France is"
trace = tl.trace(gpt2, text) # run model as-is, logging every operation
op = trace["tanh_6"] # fetch the 6th tanh operation
op.out.shape # torch.Size([1, 5, 3072])
op.parents, op.children # the operations before and after tanh_6
mlp = trace.modules["transformer.h.5.mlp"] # look up a module's record by its address
mlp.total_flops_forward # MLP's forward FLOPs (one of 800+ metadata fields in TorchLens)
trace.draw(module="transformer.h.5.mlp.act") # draw the graph of the MLP's GELU activation module
This figure zooms in on the MLP's GELU module; here is the whole GPT-2 graph. Each node is one tensor operation. For instance, tanh_6_266 is the 6th tanh operation and the 266th operation overall. The function name and call number are enough on their own, so trace["tanh_6"] and trace["tanh_6_266"] return the same operation. A trace can also be indexed by a module's address (trace["transformer.h.5.mlp"]) or by a selector (below).
That one tl.trace call runs your model as usual and logs the result of every operation it executes, not just module outputs. Because the Trace object returned by this function holds every step of the forward pass, it can answer questions that require the whole computational graph: which units influence which other units (receptive and projective fields), whether replaying each step reproduces its saved result (tl.validate), or how results change downstream when you edit one activation. A trace with gaps (e.g., module outputs only) cannot fully answer such questions.
Since the trace holds everything, TorchLens doesn't need a separate tool for every question. A small "basis set" of composable tools covers just about anything you'd want to do when inspecting or intervening on a network:
Select activations or gradients by any criteria, from simple ("every LayerNorm in the transformer.h.5 module") to sophisticated ("every unit in layer1 and layer2 of a ResNet affected by a specific set of units in conv1"). Any set of units you can describe by their kind, location and causal links to other units can be picked out with a single selector like those below (provably).
trace[tl.func("layer_norm") & tl.in_module("transformer.h.5")] # by function and module
trace["conv1"].projective_field.of(units, within=["layer1", "layer2"]) # by downstream influence (ResNet-18)Extract activations at whatever speed or level of detail you need: save them all, stop the forward pass once you have the ones you want, or quickly run a dynamic model many times while saving only those matching your criteria (e.g., every ReLU after a convolution). Saved activations can stay connected to the network if you want, so you can train it on any loss computed from them. You can also apply a function to each activation as it is saved (e.g., averaging across token positions or moving it to the CPU), and even go below Python function calls to the results of PyTorch's low-level ATen operations.
trace = tl.trace(
gpt2, text,
save=tl.func("tanh"), # keep only the tanh ops
stop_after="transformer.h.5", # stop forward pass after block 5
capture=tl.options.CaptureOptions(
aten_ops="values", # save ATen results too
activation_transform=lambda x: x.mean(1), # also save each activation averaged across token positions
),
)
trace["tanh_6"].out.norm().backward() # train on any function of saved activations
trace["tanh_6"].transformed_out # the saved average across token positions, shape (1, 3072)
trace["tanh_6"].aten_ops[0].out # output of the first ATen call beneath tanh_6Gradients work the same way: ask for them when you call tl.trace, run a backward pass from any loss you want, and read the gradient at any operation, or at the grad_fn nodes PyTorch runs between operations.
trace = tl.trace(
gpt2, text,
capture=tl.options.CaptureOptions(save_grads=True), # record gradients when a backward pass runs
)
loss = trace["lm_head"].out[0, -1].max() # any scalar can serve as the loss; here, the top next-word score
trace.backward(loss) # run the backward pass from the loss
trace["tanh_6"].grad.shape # torch.Size([1, 5, 3072]): the gradient at tanh_6's output
trace["tanh_6"].grad_fn.calls[0].grad_inputs[0].shape # torch.Size([1, 5, 3072]): the gradient passed back to tanh_6's inputInspect every part of a network in exhaustive detail. Every operation, module, parameter, buffer, and grad_fn node gets its own data structure, as does the trace itself, with 800+ metadata attributes in all, covering everything from each operation's output shape and FLOPs to the line of code that ran it, whether it was inside an if branch, and the random number state when it ran. You can summarize this metadata in a readable table, or export it to pandas (one row per operation, module, parameter, etc.).
print(trace["tanh_6"]) # everything TorchLens knows about one operation
# Layer tanh_6: float32 [1, 5, 3072] on cpu, min -1.000 max 0.9937 mean -0.4892 sd 0.4728
# position operation 266 of 541
# module transformer.h.5.mlp.act
# output 15,360 values, 60.0 KiB
# values min -1.000 max 0.9937 mean -0.4892 sd 0.4728
# graph mul_23 -> tanh_6 -> add_25 (reaches 13 of 25 model outputs)
# function tanh, gradient TanhBackward0, 151 us
# source transformers/activations.py:66 in forward
# more .ops .parents .children .to_pandas()
trace["tanh_6"].code_context[-1].call_line # the exact line of source code that called tanh_6
# 'return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0))))'
trace.params["lm_head.weight"].all_module_addresses # GPT-2 reuses its embedding weights as its output layer
# ['transformer.wte', 'lm_head']Visualize a model's computational graph using one-line defaults, premade styles for common tasks (e.g., debugging, profiling, seeing what an intervention changed), or full customization. TorchLens can show recurrent networks "rolled" or "unrolled", collapse nested modules to any depth (or pick a readable depth automatically), and draw the backward pass. Each node (one operation) can also show its own content, like input images or text, output category labels, or MDS plots of intermediate layer representations. Any visual attribute, such as node size or color, can be mapped to any metadata you like (e.g., memory use or execution time). The visualization guide explains every part of a drawing and every option.
trace.draw(collapse="auto") # fold nested modules into boxes as needed to keep the graph readable
trace.draw(rolled=True) # draw a loop's layers once each, with an arrow back for the repeats
trace.draw(color_by="time", size_by="bytes") # color nodes by execution time, size them by memory useIntervene on any activation or gradient (using simple syntax based on Judea Pearl's do-calculus) to ablate, steer, patch in activations from another run (20+ built-in edits), or apply your own function. Edits can propagate through the saved graph without rerunning the model, or be applied to a fresh run.
site = tl.module("transformer.h.5.mlp") # select block 5's MLP to edit
ablated = trace.fork().do(site, tl.zero_ablate()) # copy the trace and zero the MLP in the copy
italy = tl.trace(gpt2, "The capital of Italy is") # trace the same model on an Italy prompt
patched = italy.fork().do(site, tl.patch_from(trace)) # patch in the MLP's activations from the France runCompare any set of traces (e.g., clean vs. patched, different inputs, different checkpoints) operation by operation by grouping them into a Bundle object: compute metrics across them at any operation, find where they diverge most, and show the differences on the graph.
runs = tl.Bundle({"italy": italy, "patched": patched}) # group two Traces into one Bundle
runs.rank_changes("italy", k=10) # the 10 operations the patch changed most
runs.draw(graph="diff") # draw the model with those changes highlightedDebug models with over 20 convenient functions for finicky tasks like locating the first operation that produces a NaN, finding dead units and exploding gradients, and inferring what input shape a model expects. Each result names the operation and the line of code that ran it.
tl.debug.find_nan(model, x) # finds the first operation that produced a NaN and its line of code
tl.debug.dead_neurons(trace) # dead units in each layer
tl.debug.grad_health(trace) # vanishing and exploding gradients, after trace.backward(loss)
tl.debug.guess_input_shape(model).shape # (1, 16): input shape inferred from the model aloneEvery tool above is designed to work on any PyTorch model, including ones with loops, conditional branches, and other dynamic behavior. For example, this small model has a conditional branch and a loop. Setting rolled=True draws each operation in the loop as one node, no matter how many times it runs; the default (unrolled) gives each pass its own node:
class Loop(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(4, 4)
self.fc2 = nn.Linear(4, 4)
def forward(self, x):
if x.sum() < 0: # branch
x = -x
for _ in range(2): # 2 passes
x = self.fc2(torch.relu(self.fc1(x)))
return x
loop = Loop()
x = torch.randn(1, 4)
trace = tl.trace(loop, x)
trace.draw(rolled=True, direction="leftright") # rolled (first picture)
trace.draw(direction="leftright") # unrolled (second picture)
TorchLens also works at every scale, from single ATen calls up to agent harnesses that call the model many times. For autoregressive language models, TorchLens can trace the whole generation loop. Give it gpt2.generate, and each new token gets its own pass through the model, so you can pull out the activations for any token (guide):
gen = tl.trace(gpt2.generate, "The capital of France is", input_kwargs={"max_new_tokens": 4})
gen["transformer.h.5.mlp"].num_calls # 4: one call per generated token
gen["transformer.h.5.mlp:1"].out.shape # torch.Size([1, 5, 768]): the first call reads the whole prompt
gen["transformer.h.5.mlp:4"].out.shape # torch.Size([1, 1, 768]): later calls read only the newest tokenPipelines and agent harnesses are the next scale up, sending the model several separate requests (e.g., draft and critique, voting, tool calls), each a generation like the one above. Put one trace per request in a Bundle, tell it which answer went into which prompt, and TorchLens pairs each token of a later input with the token of an earlier output it was copied from, so you can follow an effect across calls or draw the pipeline as a graph (notebook).
For popular networks (e.g., Hugging Face models, torchvision models), TorchLens adds quality-of-life features like automatically preprocessing text or image inputs and showing ImageNet category labels for outputs. Additionally, it can split operation outputs into functionally distinct "facets" like the individual attention heads in a transformer, which you can select, inspect, and intervene on just as you would any other activation. You can add such functionality to models of your choosing, and TorchLens will add it for popular new models as they come out. For example, given a vision transformer from Hugging Face and a raw image, TorchLens applies the preprocessing published with the model, labels the outputs with ImageNet categories, and splits each attention layer into heads you can inspect or edit one at a time:
from transformers import ViTForImageClassification
vit = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
cat = tl.trace(vit, "cat.jpg") # TorchLens recognizes this model, so it preprocesses the raw image for you
cat.output_table(k=3) # tiger cat 0.44, tabby 0.25, Egyptian cat 0.17
cat.modules["vit.layers.5.attention"].facets.q.shape # torch.Size([1, 197, 12, 64]): queries split into 12 heads
head3 = tl.module("vit.layers.5.attention") & tl.head(3) # head 3 of block 5's attention
ablated = cat.fork().do(head3, tl.zero_ablate()) # copy the run and zero that head in the copy
ablated.output_table(k=3) # tiger cat 0.47, tabby 0.27, Egyptian cat 0.18The tools are built to combine, so whole experiments can fit in just a few lines:
france = tl.trace(gpt2, "The capital of France is") # for Hugging Face models, TorchLens tokenizes input text automatically
italy = tl.trace(gpt2, "The capital of Italy is")
def paris_vs_rome(trace): # how much more the model predicts " Paris" than " Rome"
logits = trace["lm_head"].out[0, -1] # next-word scores
return (logits[6342] - logits[10598]).item() # GPT-2's token ids for " Paris" and " Rome"
heads = {(b, h): tl.module(f"transformer.h.{b}.attn") & tl.head(h, "v")
for b in range(12) for h in range(12)} # values of all 144 attention heads
effects = {bh: paris_vs_rome(italy.fork().do(head, tl.patch_from(france)))
for bh, head in heads.items()} # patch each head from France into Italy
max(effects, key=effects.get) # (9, 8): the head that matters most
# patch heads (9, 8) and (10, 0) from France into a copy of the Italy run
both = italy.fork().do(heads[9, 8] | heads[10, 0], tl.patch_from(france))
paris_vs_rome(italy), paris_vs_rome(both) # (-5.25, 2.99): the model now prefers Paris over RomeFor a fuller walkthrough, see the activation patching notebook.
You can save a trace with trace.save and load it later (even on another machine) with tl.load, so you can share a result, compare two runs, or hand one to an agent without rerunning the model.
TorchLens (as you might expect) focuses on PyTorch, whose eager execution and dynamic autograd graph make every forward and backward step visible and editable mid-run. But nothing in its design depends on PyTorch itself (since all networks are just sequences of tensor operations with trainable parameters, packaged into modules—exactly what TorchLens records), so it also has backends for JAX, TensorFlow, MLX, PaddlePaddle and tinygrad (feature comparison), and future frameworks can be added the same way.
The work of TorchLens stops at the edge of the network: it gets you the activations, gradients and metadata from inside it, but leaves the downstream analyses (such as decoding, representational similarity, and dimensionality reduction) to you and to the many excellent packages built for them (see Related tools).
pip install torchlens # with pip
conda install -c conda-forge torchlens # or with conda
# Install Graphviz to draw a model's computational graph (https://graphviz.org/download/):
sudo apt install graphviz # Debian/Ubuntu
brew install graphviz # macOSTorchLens runs on Python 3.10 to 3.13 with PyTorch 2.1 or later. Features that require other packages are optional extras; for example, pip install "torchlens[hf]" allows you to pass raw text or images straight to Hugging Face and timm models (the full list is in pyproject.toml).
- Tutorial notebooks: eleven short notebooks that build up TorchLens step by step, from a first trace to running full experiments at scale. Each runs on a laptop CPU in about a minute, or in your browser on Colab.
- Guides by field: where to start, depending on your work:
- Examples: runnable notebooks and scripts, from quick tasks to full experiments:
- Recipes, one short self-contained script per common task: extract activations, find the operation that first produces a NaN, save and load a trace.
- Workflows: ablation grids, activation patching, steering, representational similarity for neuroscience.
- Replications, scripts that re-test published interpretability findings on GPT-2: induction heads, the IOI name movers, causal tracing.
- Recent methods, run end to end on public checkpoints: the Jacobian lens, EAP-IG circuit discovery on GPT-2's IOI task, attribution graphs with transcoders, the logit lens on ViT and CLIP, LEACE concept erasure, DAS causal bases.
- Glossary: definitions for every public function and argument.
- API map: an interactive, searchable map of every public module, class, function, method and argument, rebuilt from the code at each release.
- FAQ: common questions, including what TorchLens deliberately leaves to other tools.
- Also in TorchLens: 58 features that didn't fit in this README, like:
- tracing a model with no weights loaded
- knowing which field of a model's input or output each tensor belongs to, even in nested objects like a Hugging Face
CausalLMOutputWithPast(containers) - finding the first operation where fp16 or bf16 drifts from a float64 rerun (
tl.debug.bisect_precision)
TorchLens is meant to work on any model, so it is continually tested on a growing Model Menagerie of over 6,400 networks, spanning over eight decades of neural network history and every major domain (vision, language, audio, graphs, reinforcement learning, and more). Each model is checked with tl.validate, which re-executes every saved operation and confirms that the results match the model's own output. More than 6,300 models have been verified in this way so far, with more being validated all the time.
Beyond stress-testing TorchLens, the Menagerie aims to be a browsable gallery of every neural network architecture ever publicly described, from McCulloch and Pitts's 1943 network of threshold neurons (the earliest on record) to today's models. Each page carries the network's full computational graph. TorchLens draws it from a real forward pass, so it shows every operation the network performs, in the same notation for every architecture since 1943.
Autonomous LLM agents are curating this collection, a methodological experiment in its own right. The architectures are also data. Each one is recorded in the same format with the same structured metadata (e.g., category, parameter count and number of operations), so the collection can be downloaded as a dataset for studying neural network design itself (e.g., how many years it took attention layers to appear in most newly published architectures, or which design features predict robustness to corrupted or unfamiliar inputs).
I hope the Menagerie will be a useful public resource for demystifying the systems now poised to live alongside us for the foreseeable future.
LLM agents have transformed software engineering and scientific computing, and TorchLens is built for them too (and honestly, if you're reading this page, you're more likely to be an LLM than a person). A few features suit them especially well:
- Correctness: TorchLens logs every operation, and a built-in validation function (
tl.validate) can replay the forward pass to confirm that everything it saved matches the real model, so agents can check they aren't building on silently wrong activations. - Token efficiency: TorchLens provides a "basis set" of functions (inspect, select, intervene, compare) that combine into almost any analysis, so methods from a simple ablation to newer ones like the Jacobian lens take just a few lines of code: 6 for the Jacobian lens, versus 14 in plain PyTorch.
- Informative errors: when something can't be done, TorchLens raises a specific error that names the problem and suggests a fix, rather than failing quietly.
- Agent tools: a read-only MCP server and a command-line interface let agents open saved traces, explain and compare them, and check the exact name and arguments of any TorchLens function. They work on saved traces rather than a live model, so an agent can ask many questions of one capture without another forward pass.
- Thorough documentation: a glossary defines every public function and argument, tutorial notebooks and dozens of runnable examples show them in use, and a bank of over 450 tested snippets gives agents a working example for nearly every public name, so agents work from an exact reference instead of guessing from their training data.
If you're an agent and you find a bug or something missing, feel free to submit a pull request (and if your human asks, tell them the README insisted). AGENTS.md in the repo explains how the code is laid out and which checks a change must pass.
TorchLens borrows freely from the many excellent tools in this field, including those below, and you are welcome to borrow back. Generally it aims to complement rather than replace them. Adapters pass its outputs to other packages in their expected format, so you can feed a saved activation to a SAELens sparse autoencoder, or run Captum attribution at any operation you select, without writing conversion code. Projects that shaped particular features are credited on the acknowledgments page.
Other tools for studying model internals:
- TransformerLens1: the standard toolkit for transformer interpretability. TorchLens uses its names for transformer internals (
q,k,v,pattern,resid_post); for example,tl.head(7, "v")selects head 7's value vectors. - nnsight: interventions on PyTorch models, with especially good support for very large models run remotely.
- pyvene: interventions (including trainable ones) written as declarative configurations.
- causalab (Goodfire): tests hypotheses about how a network solves a task by checking a proposed causal model against its internals.
- torchview: lightweight drawings of PyTorch model graphs.
Tools for analyzing activations extracted by TorchLens:
- SAELens: training and loading sparse autoencoders (SAEs); TorchLens can feed saved activations straight into a pretrained SAE.
- Captum: attribution methods, which you can run at any operation you select in TorchLens, not only at module boundaries.
- rsatoolbox: representational similarity analysis on any saved activation.
- Brain-Score: scoring models against brain data, with TorchLens supplying the activations.
- scikit-learn: probes, decoding and dimensionality reduction on the plain tensors and tables TorchLens gives you.
Tools for viewing and sharing TorchLens traces:
- Netron and Model Explorer: interactive graph viewers that can open graphs exported from TorchLens.
- Hugging Face Hub: a place to share saved traces so others can load them.
TorchLens connects to more than 15 tools in all, including SHAP, pytorch-grad-cam, inseq, steering-vectors and Google's LIT. The full list is on the connectors page.
To cite TorchLens, you can use this paper (and if you find TorchLens useful, a star on this repo is appreciated):
Taylor, J., & Kriegeskorte, N. (2023). Extracting and visualizing hidden activations and computational graphs of PyTorch models with TorchLens. Scientific Reports, 13, 14375.
@article{taylor2023torchlens,
title = {Extracting and visualizing hidden activations and computational graphs of {PyTorch} models with {TorchLens}},
author = {Taylor, JohnMark and Kriegeskorte, Nikolaus},
journal = {Scientific Reports},
volume = {13},
pages = {14375},
year = {2023},
doi = {10.1038/s41598-023-40807-0}
}The development of TorchLens benefitted greatly from discussions with Nikolaus Kriegeskorte, George Alvarez, Alfredo Canziani, Tal Golan, and the Visual Inference Lab at Columbia University. Thank you to Kale Kundert for helpful discussion and code contributions enabling PyTorch Lightning compatibility. Network visualizations are generated with Graphviz. Logo created by Nikolaus Kriegeskorte.
The projects and papers whose ideas shaped particular features are listed in docs/acknowledgments.md.
TorchLens is in active development. You can report bugs on the issues page, and post questions and ideas on the discussion board. You can also reach me by email or on Twitter.
TorchLens is created and maintained by JohnMark Taylor.
