Skip to content

Also in TorchLens

The README shows the core of TorchLens. This page lists the rest, grouped by task, one line per feature, each linked to the page that documents it. Features whose names may still change are described in words only; the linked page gives the current spelling.

Capture

  • Records the dict, tuple, list and dataclass structure of a model's inputs and outputs (capture_container_structure=True). Containers
  • Lets you register your own container types so TorchLens can walk and rebuild them (tl.types.register_container). Containers
  • Saves each op's random state when save_rng=True is set, so replays redraw the same dropout masks. Glossary
  • Records each version of a registered buffer (for example BatchNorm running stats) as a graph node (Op.graph.is_buffer). Buffers
  • Records torch.func transforms such as vmap and grad as single boundary ops (Op.call.is_func_transform). Glossary
  • Traces a model built on the meta device, with no weights loaded, and gets its full graph, module tree and every op's shape and dtype. Weights-free capture
  • Captures a multi-step loop that calls one model many times as one episode. Episode capture
  • Keeps each generated token's new KV-cache slices for the decoder layers you choose. KV cache
  • Streams selected activations to a bundle on disk during capture instead of holding them in memory. Performance
  • Merges the captures from each rank of a distributed run into one view, joined at their collective ops (tl.merge_ranks). Glossary
  • Extracts chosen layers over a whole dataset in batches to disk, and picks up where it left off after an interruption. Extract features
  • Records an event stream instead of a graph, and builds a full trace from it on request (tl.record). Glossary
  • Recovers a recording that a crash or interrupted run left half-written (tl.recording.recover). Glossary

Reading

  • Runs a dataloader through a model and accumulates running means, covariances, PCA or CKA of any site's activations or gradients (tl.stats.aggregate). Streaming statistics
  • Prints a live stats line for a site each time it runs. Value displays
  • Selects units that are dead, always saturated, or barely varying across many inputs. Intervention API
  • Selects units by how much they vary across a recurrent layer's repeated passes. Intervention API
  • Selects every op on the paths between two regions, or within a given number of hops of one. Intervention API
  • Draws a seeded random selection of the same size as any selection, as a control. Intervention API
  • Selects the part of an activation along a direction or subspace, such as a probe vector. Intervention API
  • Checks computed receptive fields against gradients from the real model (tl.receptive_field.verify). Receptive and projective fields
  • Builds representational dissimilarity matrices (RDMs) and compares them. Neuro reference
  • Computes linear CKA with a bootstrap confidence interval over stimuli. Glossary
  • Shrinks activations with a seeded sparse random projection as they are saved. Activation transforms
  • Writes a plain-language report of what a trace contains. Report surfaces
  • Tables time and memory per op, module or call (trace.profile()). Report surfaces

Editing

  • Replaces a site's output with the output of a module you supply (tl.splice_module). Glossary
  • Removes, or keeps only, an activation's component along a direction (tl.project_out, tl.project_onto). Glossary
  • Limits an edit to chosen token positions, including each row's first or last real token under padding. Glossary
  • Packages edits as one reusable object and saves it to a file. Intervention API
  • Runs one edited trace per value and returns them together as a bundle (tl.sweep). Glossary
  • Sweeps every combination of sites and edits into a labelled grid of results. Experiments
  • Patches one attention head, or its attention pattern, from a clean run into a corrupted run, exactly, even when the model uses fused attention (SDPA). Facets

Gradients and training

  • Gives each backward pass its own record, and selects ops by backward pass number (tl.in_backward_pass). Glossary
  • Selects the gradients flowing into or out of each autograd step (tl.grad_input, tl.grad_output). Glossary
  • Clips the gradient at any op to a maximum norm during the backward pass (tl.clip_norm). Glossary
  • Finds where NaN or Inf gradients first appeared in a backward pass. Debug reference
  • Finds the first op whose output drifts from a float64 rerun of the same model by more than its dtype explains. Debug reference
  • Runs training-loop checks on a model and optimizer without any capture (tl.debug.ChecksSession). Training checks
  • Warns when mixed-precision training never accepts an optimizer step (tl.debug.Watchdog). Training checks
  • Streams gradient and activation statistics from a training run to TensorBoard, W&B, MLflow or Comet, with callbacks for Lightning and the Hugging Face Trainer. Trackers

Drawing

  • Draws every layer and head's attention pattern on one overview page, in the style of BertViz's model view, with a detail page per layer. Transformer pictures
  • Animates the repeated passes at one site of a recurrent model (trace.animate). Glossary
  • Shows activation heatmaps or channel grids inside the graph's nodes (tl.viz.feature_maps). Glossary
  • Draws how a site's value distribution changes over training steps. Training ribbon

Saving and sharing

  • Exports per-op timing as Chrome tracing, speedscope or flamegraph files (tl.export.chrome_trace). Export reference
  • Exports a timeline of tensor memory over the forward pass (tl.export.memory_timeline). Export reference
  • Exports saved outputs as an xarray Dataset in Brain-Score's layout (tl.export.xarray). Export reference
  • Writes a single-file HTML report of a trace that opens offline (tl.export.html). Offline report
  • Exports a trace's tensors to one safetensors file and reads them back. Export reference
  • Fails if a model's structure differs from a pinned hash. Fingerprint reference

Tooling

  • Checks your environment and model for known trouble spots, such as offloaded or quantized layers (tl.report.doctor). Report surfaces
  • Gives code written against torchextractor a drop-in Extractor class (tl.compat.Extractor). Coming from torchextractor
  • Registers your own output decoder, such as class names for a custom classifier (tl.postprocessing.register). Glossary
  • Accepts TransformerLens names like hook_pattern for facets, if you turn it on (tl.facets.enable_transformer_lens_aliases). Facets
  • Renders traces as Treescope cards in notebooks. Treescope cards
  • Raises named error types with specific codes and remedies (tl.errors), so code can catch one failure and tell it from the others. Capture outcomes
  • Ships a Claude Code plugin with a TorchLens skill that tells agents which call to use. For AI agents