Skip to content

Intervention Visibility Classes

TorchLens can intervene on operations that its decoration layer sees during a normal Python forward pass. In practice, that means tensor operations routed through wrapped PyTorch callables, plus module output boundaries that TorchLens records during capture.

Leaf Op

A leaf op is a single decorated PyTorch operation such as torch.relu, torch.add, torch.matmul, or a tensor method that TorchLens wraps. These are the most direct intervention targets:

log = tl.trace(
    model,
    x,
    save=tl.func("relu"),
    intervene=tl.when(tl.func("relu"), tl.zero_ablate()),
)

For post-hoc replay workflows, capture enough payloads and replay metadata, then fork and attach hooks after discovery. Capture-time intervene= is the direct single-forward path when the site is known in advance.

Compound Module

A compound module is a user or library nn.Module that contains several leaf ops. TorchLens records module entry and exit metadata, so you can discover inside a module with tl.in_module("encoder.layers.0") and then target a finer leaf op, or target a module output boundary with tl.module(...) when that is the right semantic site.

The conservative workflow is discover first:

sites = log.find_sites(tl.in_module("block"), max_fanout=32)
target = tl.label(sites.where(lambda site: site.func_name == "relu").labels()[0])

Opaque Fused Kernel

Opaque fused kernels execute internal math that is not exposed as separate Python-level PyTorch calls. TorchLens can target the visible fused call, but not hidden intermediates inside it.

Important examples:

  • torch.nn.functional.scaled_dot_product_attention is fused from TorchLens' perspective. You can intervene on the visible SDPA call output, but not the internal Q/K/V scores, attention probabilities, or value aggregation steps.
  • FlashAttention kernels have the same limitation: the kernel internals are not TorchLens leaf ops. A FlashAttention 4 call (flash_attn.cute) is one autogradfunction boundary op whose output you can intervene on.
  • Multi-output fused calls, including nn.LSTM, expose one visible TorchLens op per returned tensor such as sequence output, final hidden state, and final cell state. The fused kernel is still opaque internally.
  • Compiler stacks such as torch.compile, TorchScript, and exported graphs can bypass the Python calls TorchLens decorates. Capture the original uncompiled nn.Module when you need intervention visibility.

If you need intervention access to Q/K/V intermediates, use a manual attention implementation that calls visible PyTorch operations for projections, scores, softmax, and value aggregation.