TorchLens: tracing every activation and gradient in a PyTorch model

The problem TorchLens addresses: seeing what a model actually computed

Reading a training curve tells you that a loss went down. It does not tell you what happened inside the network, which tensor became the input to the layer you care about, whether a gradient reached a parameter, or how the graph was shaped when a dynamic control flow took a different branch. The usual answer is a forward hook, and the usual complaint is that hooks miss things.

TorchLens takes that complaint seriously enough to make it part of the pitch. The README describes the package as capturing every activation and gradient across the forward and backward pass, auto-visualizing the full computational graph, exposing rich per-op metadata, and allowing intervention on the network as it runs, with the claim that this works on any architecture, including dynamic and recurrent ones.

Two phrases there are worth separating. The first is the capture itself: one call returns a log object holding the graph, the activations and the metadata, rather than a set of ad hoc references you assembled yourself. The second is the auto-visualization: the same log can be drawn as a PDF of the computational graph, which turns the capture into something you can put in a review, a bug report or a lecture.

The intended audience is research and engineering work where the model's internals are the object of study. The repository's own descriptions line up with that: the package summary in pyproject.toml calls it a package for extracting activations from PyTorch models and saving portable bundles, and the keyword list is torch, torchlens, activations, features and io. The project also points at a paper in Nature Scientific Reports, a ten-minute tutorial notebook, and separate five-minute and fifty-minute galleries, which suggests a tool built by someone who expects other people to read its output.

The trace call and the three-rung input ladder

The central verb is tl.trace, and the README frames the ways of calling it as a ladder with three rungs, from most to least faithful:

# API sketch; `model` is your model and `lm` any HuggingFace language model.
log = tl.trace(model, x)                             # best: your real input
log = tl.trace(model, input_size=(1, 3, 224, 224))   # your shape, random values; disclosed
log = tl.trace(model)                                # inferred shape + random values; disclosed, or a teach
log = tl.trace(lm, "The quick brown fox")            # HF models: a string is a real input

The design point is disclosure. The second and third rungs do not know what your model will actually see, so they substitute random values, and the README marks those calls as disclosed in the code comments themselves. For a dynamic model whose behaviour depends on input content, that distinction decides whether your trace describes the model or a plausible-looking stranger. The fourth rung handles HuggingFace language models, where a string is a genuine input rather than a shape placeholder.

A worked example follows the same logic. The README's snippet loads a pretrained ResNet-18 and calls .eval() before tracing, with a comment explaining why: pretrained models arrive in train mode, and a train-mode forward silently updates BatchNorm running statistics. That detail is the kind of thing a hook-based workflow usually discovers the hard way, after the fact.

import torch, torchvision.models as models, torchlens as tl

model = models.resnet18(weights="IMAGENET1K_V1").eval()  # .eval(): pretrained models arrive
# in train mode, and a train-mode forward silently updates BatchNorm running statistics
x = torch.rand(1, 3, 224, 224, generator=torch.Generator().manual_seed(0))

log = tl.trace(model, x)      # one call -- full graph + all activations
print(log.summary())          # module table, op count, FLOPs
print(log["conv2d_1_1"].out.shape)  # grab any activation by name ...
print(log["layer1.0"].out.shape)    # ... or by module path
print(log[7].func_name)             # ... or by ordinal
log.draw()                    # PDF of the computational graph

That snippet is also a compact tour of the log API. summary() returns a module table with an op count and FLOPs. Individual tensors are reachable by generated name, by module path, or by ordinal position in the graph, which is the practical escape hatch when you do not know how an op was named.

Installing TorchLens and the Graphviz dependency

The installation order in the README is explicit: install Graphviz first, because it is required for graph visualizations, then install TorchLens:

sudo apt install graphviz   # Debian/Ubuntu; see graphviz.org for other platforms
pip install torchlens

The ordering matters because the failure mode is late. If Graphviz is missing, the package still imports and tracing still works, and you find out at the point where you call log.draw() and expect a PDF. Installing it first removes that ambiguity, and the comment points at graphviz.org for platforms other than Debian and Ubuntu.

The supported version matrix is stated in a table rather than a sentence, and the README makes a point about how that table is maintained: each row is executed by a CI leg on every release, not asserted from metadata. Python is supported on 3.10, 3.11, 3.12 and 3.13, with one smoke-tier CI row per interpreter plus clean-venv wheel installs at release. PyTorch support is more conditional: 2.1 and above on Python 3.10 and 3.11, 2.2 and above on 3.12, and 2.6 and above on 3.13, the first torch release with wheels for that interpreter. Floor rows pin the exact minimum pair per interpreter, newest-admitted rows pin the current release, and a daily canary runs against the latest.

The pyproject.toml backs the Python floor with a mechanical reason: the package uses dataclass slots in 143 places across 31 modules, and slots requires Python 3.10, so import torchlens cannot succeed on 3.9. The build system also pins setuptools at version 83 or newer, with a comment explaining that older versions mishandled MANIFEST.in exclude and prune directives without Unicode normalisation, which on macOS filesystems could let an NFD-named file slip past an NFC exclusion and be packed into the source distribution.

The runtime dependencies include numpy, packaging, tqdm and graphviz, and the comment in the manifest is explicit that those floors are declared and exercised, so the tested minimum and the declared minimum are the same artifact.

What a capture actually contains: op-level metadata and gradients

The volume of metadata is the headline claim. The README states that TorchLens records 180 or more metadata fields per operation, and 550 or more fields in total across every record type, which the list of record types enumerates as operations, modules, parameters, buffers, gradients, and the model itself.

That inventory explains why the tool is more than an activation dumper. Parameters and buffers are the things a model carries between calls, gradients are what the backward pass produces, and the model record holds the whole picture. If you are diagnosing a training bug where a gradient never arrived, having the gradient records in the same object as the activations and the graph is the point.

The README also contrasts this with what plain hooks miss, and the list is specific rather than generic: dynamic paths, reused modules, functional ops, fused attention, recurrent unrolls, or the metadata needed to tell two sites apart. Functional ops are the clearest case, because a module hook fires on nn.Module boundaries while a functional call such as a fused attention op has no module to hang a hook on at all. Reused modules raise a different problem, since one hook on a shared module cannot distinguish the two call sites without extra bookkeeping.

Live intervention is the third pillar, and it is the one the README summarises least, beyond the statement that you can intervene on the network as it runs. The repository shows where that lives: there is an examples/intervention/ directory, alongside examples/recipes/, examples/demos/, examples/debug_utilities.py and examples/runnable_tlspec.py, which suggests intervention is a supported surface with worked examples rather than an incidental feature. The documentation index points to a performance guide, a page on receptive and projective fields, an AI-agent quick reference, and a limitations and remedies page.

How the faithfulness claim is checked: replay, parity and tripwires

A capture tool that silently misses an op is worse than no tool, because you will trust the picture. TorchLens answers this with a validation approach described at length in the README, and the claim is hedged in a way that deserves attention.

The mechanism is a pipeline:

model forward -> TorchLens capture -> DAG replay -> output parity + metadata invariants

The captured DAG is replayed op by op and its outputs are compared against the original forward pass, with metadata-invariant tripwires run over the resulting graph. If a replay or an invariant fails, the README is emphatic, the capture is treated as genuinely wrong and caught automatically, not waved through because the model ran.

The validation campaign is run against a catalog of architecture entries, which the README describes as 8,500 or more rows across roughly 3,600 families, measured on 2026-08-16, spanning image, video, audio, multimodal and language models across feedforward, recurrent, transformer, GNN, MoE and diffusion families. The wording hedges: the public claim stays hedged until the campaign closes, and the exact per-row verification status lives in the menagerie project's append-only ledger rather than in the README.

The residual is reported rather than hidden. The remaining tail is described as mostly genuinely huge models, models with hard-to-build dependencies, or architectures that need more bespoke inputs before an automated run is fair, and those are counted as unverified rather than presented as successes. Inside the repository itself, a coverage-chosen sample of hand-built classics under tests/classics_corpus/ runs the same capture, replay and invariant checks as part of the test suite.

So there are two confidence levels, and the difference matters for how much weight you put on a specific capture. A model in the classics corpus is checked by the test suite. A model in the menagerie campaign is checked by the campaign and recorded in an external ledger. Neither tells you that your particular checkpoint behaves, which is why the capture itself is worth running on your own inputs.

Getting the graph out, and what the log lets you do with it

Drawing the graph is a method on the log rather than a separate tool: log.draw() produces a PDF of the computational graph, and Graphviz is the reason it is a dependency rather than an optional extra. For a model with hundreds of operations, that PDF is often the artefact that gets shared, because it shows shape and connectivity at a glance in a way a tensor dump never will.

The rest of the API is deliberately low ceremony. log.summary() gives a module table, an op count and FLOPs, which is the first thing to look at when you want the shape of the network without opening the drawing. Tensors are then addressed directly by three different keys: a generated operation name such as conv2d_1_1, a module path such as layer1.0, or an ordinal index such as 7. The ordinal route matters most in practice, because op names are generated and may not match what you would have guessed, while an index into the recorded sequence always refers to something real.

The repository also documents portable output. The package description mentions saving portable bundles, and the keyword list ends with io, and the tree contains examples/semantic_io_legibility_demo.py and examples/treescope_cards_gallery.py, plus a scripts/ directory, templates/ and tools/. A capture that can be written out and reopened is what makes a graph comparable across two runs or shareable in a bug report without shipping the model.

Several migration and comparison documents sit at the top of the repository, including MIGRATIONS.md, LIMITATIONS.md, RESULTS.md, ROADMAP.md, CHANGELOG.md and a docs/migration/ directory with migration tables. Combined with the Limitations and remedies page in the docs and the AGENTS.md and CLAUDE.md files at the root, that is a project that documents its edges deliberately.

Limits, costs and the cases where plain hooks are still enough

The obvious limitation is the one the README itself lists first among its misses: plain forward hooks can be fast and useful. They are also a couple of lines of code. For a model where you want one tensor from one layer, a hook is faster to write, has no dependency, and does not capture a full graph.

The second limitation is environment. Python 3.10 is a hard floor because of the dataclass slots usage, and PyTorch 2.1 is a floor on Python 3.10 and 3.11, with the floor rising to 2.6 on Python 3.13. A team on an older interpreter cannot install the package at all, and a team on Python 3.13 needs a recent torch. Graphviz is a third, smaller dependency that has to be present before drawing works.

The third limitation is cost during capture. A full graph capture of every activation and gradient across both passes is not free in memory or time, and the README points to a separate performance guide in docs/performance.md rather than making sweeping claims, which suggests the trade-off is real and worth reading before you turn it on in a training loop. The repository also carries a benchmarks/ directory and requirements-test-tiers.txt with test-tiers.json describing tiered test coverage.

The fourth limitation is coverage of the claims themselves. The validation campaign is described as still open, with the public claim hedged until it closes, and the residual tail of unverified models is attributed to model size, dependency complexity and input requirements. So TorchLens can tell you what it captured in your run, and it can tell you when its own replay disagrees with the forward pass, but the assurance that the tool works on your architecture family is something you confirm by capturing it yourself.

The fifth is packaging hygiene rather than functionality. The manifest pins setuptools at 83 or newer because of the MANIFEST.in normalisation issue on macOS, and the repository carries a .gitleaks.toml, .secrets.baseline and a SECURITY.md, which is a sign of a project that takes supply chain hygiene seriously and expects users to build from source distributions at least occasionally.

Editorial conclusion

TorchLens fits teams that must explain what a model computed rather than merely run it: researchers debugging architecture behaviour, engineers auditing a deployed network, and anyone who has grown tired of forward hooks that miss functional ops and dynamic paths. It does not fit a training loop that cannot afford graph capture overhead, nor a team stuck on Python 3.9 or PyTorch below 2.1. Verify first that Graphviz is installed, because log.draw() depends on it, and that the versions your environment pins fall inside the supported matrix, which the README says is exercised by CI rather than asserted from package metadata.

Frequently asked questions

Official sources

1. Issues
2. johnmarktaylor91/torchlens on GitHub
3. License: Apache-2.0
4. README
5. Releases