Backward Capture¶

TorchLens records backward execution with two clocks. The forward event stream remains the source for operations, activations, modules, and graph position. Backward execution is written to a runtime-only sidecar diary, then projected into BackwardPass, GradFn, GradFnCall, and per-op gradient records. Saved .tlspec files persist the projections and tensor blobs, not the live diary.

Backend status: torch is the only true-backward capture backend in this release. MLX and TensorFlow declare backward capture unsupported. JAX, tinygrad, and Paddle expose derived-gradient previews: leaf gradients from a second backend AD execution over declared leaves, not GradFn/GradFnCall records or live backward hooks. Real non-torch backward graphs remain research.

Capturing Gradients¶

Use capture=tl.options.CaptureOptions(save_grads=...) to choose which operation gradients are retained (capture=CaptureOptions(save_grads=...); the bare flat kwarg is removed):

trace = tl.trace(model, x, capture=tl.options.CaptureOptions(save_grads=True))
trace.backward(trace[trace.output_layers[0]].out.sum())

relu_trace = tl.trace(model, x, capture=tl.options.CaptureOptions(save_grads=tl.func("relu")))
relu_trace.backward(relu_trace[relu_trace.output_layers[0]].out.sum())

save_grads=True saves all observed op gradients. False or None observes backward structure without retaining gradient tensors. Predicate expressions use the same forward selector language, plus backward selectors such as tl.grad_fn(...) (by class name, or label= for a label substring), tl.without_op() (the former tl.intervening() alias is removed), tl.grad_input(), tl.grad_output(), and tl.in_backward_pass(k).

storage=tl.to_disk(path) streams retained gradient payloads into a .tlspec bundle. The standing trace policy can be widened or narrowed per trigger:

trace.backward(loss, save_grads=tl.func("attn") & tl.in_backward_pass(2))

Plain loss.backward(), torch.autograd.backward, torch.autograd.grad, trace.backward(...), and recording backward calls are all triggers when the autograd graph belongs to a live TorchLens trace. If an engine path bypasses those wrappers, tensor hooks open an implicit pass. Consecutive bypassed backwards may merge because hook events alone do not reveal a hard engine boundary; use TorchLens triggers for exact pass boundaries.

A Recording captures gradients only when the tl.record call asked for a backward with backward_ready=True, save_grads= or default_grad=; without one it installs no gradient hooks and Recording.backward() refuses with recording_backward_not_requested. A capture entered under torch.no_grad() or torch.inference_mode() with no backward option implies inference_only=True, and its backward() refuses naming that cause. A forward that turns grad back on itself (with torch.enable_grad():) keeps its autograd graph, so such a capture is not inference-only and backward() stays available.

Backprop To The Stimulus¶

For activation maximization, adversarial examples, and GAN inversion, the optimized stimulus can be either a plain leaf tensor or an nn.Parameter owned by an outer optimizer:

z = torch.nn.Parameter(torch.randn(1, latent_dim))
trace = tl.trace(generator, z, capture=tl.options.CaptureOptions(save_grads=True))

loss = score(trace[trace.output_layers[0]].out)
trace.backward(loss)

optimizer.step()  # reads z.grad
stimulus_grad = trace[trace.input_layers[0]].grad

TorchLens clones forward inputs before capture so it can label and move tensors without mutating the caller's object. The clone is not detached, so gradients still propagate back to the original leaf tensor or nn.Parameter; x.grad / z.grad is the optimizer-facing gradient. TorchLens also records the same input gradient on trace[trace.input_layers[0]].grad when gradient saving covers the input. If the original input was an nn.Parameter, the input op records input_was_parameter=True; it is still represented as an input node in the graph, not as a model parameter.

Records¶

trace.autograd.backward_passes is indexed positionally with Python's 0-based rules, and supports named 1-based lookup with trace.autograd.backward_passes.for_pass(k). trace.autograd.last_backward_pass returns the latest projected pass.

BackwardPass.backward_call_context stores the Python source location where trace.backward(loss) was invoked. It is a FuncCallLocation, so it exposes fields such as file, line_number, and func_name. Implicit passes inferred from orphan tensor-hook emissions have no explicit TorchLens call site and leave backward_call_context as None.

GradFn records represent distinct autograd node objects. Labels use backward-native names such as addmm_back_1_4; the forward correspondence is metadata, not part of the label. A fired node has local dense call labels like addmm_back_1_4:2. Positional access remains 0-based, while colon suffixes and pass numbers are 1-based:

grad_fn = trace.autograd.grad_fns["addmm_back_1_4"]
first_call = grad_fn.calls[0]
second_pass_call = grad_fn.calls.for_pass(2)

op.autograd.grads is a local-dense accessor over saved per-pass gradient records. op.grad returns a tensor only when exactly one gradient was saved. With zero or several saved gradients it raises a specific actionable error; use op.autograd.grads[...] or op.autograd.grad_for(bwd=k) for multi-pass traces.

param.grad remains the live accumulated PyTorch tensor. Per-pass parameter increments are only available when autograd fires AccumulateGrad, so torch.autograd.grad calls that do not accumulate into .grad do not populate parameter increment records.

Multiple Passes And Accumulation¶

Each backward engine invocation becomes one BackwardPass. Loops are not backward passes and do not affect backward naming. Retained op gradients default to all observed passes when gradient saving is enabled, so repeated backwards over a retained graph produce multiple entries in op.autograd.grads.

BackwardPass.save_grads_policy records the policy snapshot active for that pass. This makes mixed policies auditable when one trace has several backward triggers.

Higher Order¶

When create_graph=True, TorchLens re-walks newly-created autograd nodes at bracket end. Creator-attributed nodes get order = creator.order + 1; forward-created nodes are order 1. Mixed-order passes are normal because earlier-order nodes may refire while higher-order nodes are created. When attribution cannot be resolved deterministically, order is None and BackwardPass.order_attribution_coverage reports the resolved fraction.

Intervention And Replay¶

Gradient interventions address a gradient as a site: tl.grad("relu_3") is the gradient of relu_3's output, bare "relu_3" is its activation, and tl.grad("relu_3") | "relu_3" is both. Use intervene=tl.when(tl.grad(<site>), <edit>) on tl.trace, where the edit is an ordinary one (tl.scale, tl.zero_ablate, tl.clamp, tl.add_noise, tl.clip_norm) or a callable fn(grad, *, ctx). The capture arms the edit and trace.backward(loss) applies it. On an existing trace, trace.do(tl.grad(<site>), <edit>, model=model, inputs=x) reruns the forward with the edit armed (audit status "armed" until a backward settles it); pass loss=lambda t: ... to arm and run the backward in one call. Replay cannot edit gradients (replay_gradient_edit_unsupported).

Differentiable replay creates a new trace: do replays the edit onto a new fork and returns it (replay engine only; rerun and set-only ignore the option):

options = tl.options.InterventionOptions(differentiable=True)
patched = source_trace.do("relu_3", tl.scale(0.5), intervention=options)
patched.backward(patched[patched.output_layers[0]].out.sum())

Saved frontier tensors enter replay as fresh leaves, so gradients flow within the replayed computation and do not leak into the source run's buffers. Parameters upstream of the replay frontier therefore get no gradient: to train them through an edit, run the edit on a live forward (tl.when(site, edit).bind(model)(x)), and tl.intervention.keep_grad(edit) warns keep_grad_replay_frontier once per session when it fires under differentiable replay.

Backward-only replay from a saved trace is intentionally deferred. A saved TorchLens backward projection does not contain PyTorch's live autograd graph or saved-tensor closures; reconstructing that would require a TorchLens-side autograd engine. Live multi-pass capture, live backward intervention, and differentiable replay cover the useful in-session workflows.

Module Containment¶

Paired backward nodes inherit the paired forward op's module containment. AccumulateGrad nodes use the parameter's module home. Intervening in-forward nodes are inferred from neighboring op-anchored nodes when possible. Post-forward loss-construction nodes before the first op-anchored node carry no module membership. module_membership_source is "paired", "inferred", or None.

Checkpoint Invocation Witness¶

Non-reentrant torch.utils.checkpoint invocations are witnessed by capture-time TOKENS (experimental surface): the patched saved_tensors_hooks.__enter__ mints one per-trace ordinal token per classified _checkpoint_hook enter on the armed owner thread outside any engine invocation, and installs per-instance token-bearing pack/unpack wrappers. Pack evidence is count-only — the forward-side slot-to-op binding is deliberately NOT claimed (pack hooks run before TorchLens logs the producing op). Unpack evidence points are backward-derived: the grad-fn fire brackets containing them resolve through the shipped user-op pairing to L1 site-key candidates. The projected summary lives on trace.checkpoint_invocation_witness (persisted as of the tlspec v8 coordinated bump): token count, per-token pack counts / unpack window evidence / site-key candidates, degrade flags, and an evidence-scoped completeness verdict. Degrade flags cover: classifier unavailable, patch unavailable, exotic subclass, the unmatched-backward warn, the reentrant node sentinel (CheckpointFunctionBackward in the discovery stream — reentrant checkpointing is definitionally token-free), unwitnessed checkpoint enters (paused logging, non-owner thread, inside an engine invocation), and an unpreserved hook-identity attribute (torch >= 2.14's _checkpoint_internal_hook stashes a private _user_hooks attribute directly on the pack-hook callable across its __enter__/__exit__ pair; when the token wrapper cannot carry that attribute forward, the replacement is skipped instead of installed, so no token mints for that enter and torch's own hook-stack accounting is never desynced). Any flag withdraws the affirmative "no checkpoint invocation observed" verdict. The private _checkpoint_hook class resolves through the compat chokepoint behind the named HAS_CHECKPOINT_HOOK_CLASS capability flag (visible in tl.report.doctor()); a torch without it degrades fail-closed — the classifier mints NO tokens, never a false one. The sibling HAS_CHECKPOINT_INTERNAL_HOOK_CLASS flag (same visibility) tracks whether this torch runtime's checkpoint hooks carry that private identity state at all; its absence (torch 2.13 and earlier) means the token wrapper's hook-object replacement is unconditionally safe. The typed checkpoint-ambiguity refusal is a pending contract amendment and lands with the identity-read accessors once the amendment is ratified.

Executed Backward FLOPs By Node And Forward Op¶

A captured backward records every ATen call it dispatches (BackwardPass.aten_ops), and tl.report.backward_status(trace) counts them exactly for the matmul and convolution families (fma2, FlopCounterMode's convention). BackwardPass.flops_attribution charges each counted call to the GradFn call it ran under and to the forward op that node differentiates (experimental): one BackwardFlopsRow per GradFn call (grad_fn_label, call_index, op_label, flops, recompute_flops). Calls that ran under no recorded node (the engine's seed gradient, nodes no capture walked) go to the unattributed bucket, which names every call (unattributed_ops); the rows plus unattributed_flops equal the pass's executed_flops exactly. backward_status(trace).flops_attributions holds one attribution per counted pass. When one backward reaches the nodes of two captures, each capture's view keeps the nodes whose forward op it recorded and leaves the others to their capture (other_trace_ops); a pass that raised before finishing has no count.

Under activation checkpointing the backward replays forward kernels. The engine runs backward kernels with grad mode off unless the pass builds a graph, while checkpoint replays run with grad mode on (DispatchedOp.grad_enabled), so on a trace whose checkpoint witness saw checkpointing a counted call made with grad mode on is labeled recompute on the row it ran under. The label splits that row's count and never moves a FLOP between rows. A create_graph pass that checkpoints, or a trace without checkpoint evidence, reads the recompute share as None (unknown), never a guess. Matching a call to the node that ran it needs torch's executing-node resolver; on a build without it (HAS_CURRENT_AUTOGRAD_NODE is False in tl.report.doctor()) no call is charged to a node and the executed count reads unknown with that reason.

tl.report.cost_tree(trace) carries the same counts per module call: self_forward_mac_flops (the forward's multiply-accumulate work, 2 x true MACs), self_backward_flops and self_forward_backward_flops, with subtree_* twins. The total adds the backward to the forward's MAC work, one convention, never to the elementwise-inclusive self_flops. Backward work on a node with no forward op and the unattributed calls get their own rows, so the backward column sums to the executed count. Without a counted backward the columns are None and the tree's backward_state says whether the backward was not run or is unknown, never 0.

import torch
import torchlens as tl
from torchlens.report import backward_status, cost_tree

model = torch.nn.Sequential(torch.nn.Linear(4, 6), torch.nn.ReLU(), torch.nn.Linear(6, 3))
trace = tl.trace(model, torch.randn(5, 4), capture=tl.options.CaptureOptions(backward_ready=True))
trace.backward(trace[trace.output_layers[0]].out.sum())

status = backward_status(trace)
(attribution,) = status.flops_attributions
attributed = sum(row.flops for row in attribution.rows)
assert attributed + attribution.unattributed_flops == status.executed_flops
by_op = attribution.flops_by_op()  # forward op label -> backward FLOPs
tree = cost_tree(trace)
assert sum(row.self_backward_flops for row in tree.rows) == status.executed_flops
trace.cleanup()

Validation¶

Backward validation checks parameter-gradient parity, module-output gradient parity through a separate stock run with minimal tensor hooks, flow consistency for non-intervened quantities, and structural invariants. Invariants include dense pass indices, local-dense call ordinals, one record per retained autograd object, resolved order chains, label grammar, and containment flags.

The module-output check never compares the root module call, because its output gradient is the loss gradient itself. A model whose forward calls no child module therefore has no module-output gradient to compare. Validation then returns False with a TorchLensWarning (code backward_validation_no_module_output_evidence) saying so: the check has no detection power, so the result is unverifiable rather than passed. validate_layer_grads=False checks parameter gradients alone.

Dataframes And Serialization¶

trace.to_pandas() never reads loud gradient properties. Ambiguous op gradients appear as None in the grad column with num_saved_grads set. Backward tables are available through trace.autograd.grad_fns.to_pandas(), trace.autograd.grad_fn_calls.to_pandas(), and trace.autograd.backward_passes.to_pandas().

Unified .tlspec manifests include backward_summary, gradient blob kinds in body_index, and the public policy for old bundles: legacy class remaps are accepted at load, while removed backward gradient configuration fields are dropped during state normalization.

The per-call census (BackwardPass.aten_ops) is session-time. Saving writes each pass's counted executed backward and its attribution in a private persisted record on the BackwardPass, so a loaded trace reads the same backward_status count, rows and unattributed bucket.

Limitations And Costs¶

Torch keeps strong references to discovered autograd wrapper objects until trace.cleanup() or trace deletion, preventing id reuse from corrupting registry lookups. This can grow memory for large or higher-order graphs; PyTorch also warns about reference cycles with create_graph=True.

Live traces register their known autograd roots so wrapped engine calls can find owning traces. If you keep a trace only for analysis while continuing training, call trace.autograd.disarm() to detach it from engine interception and tensor-hook emission.

MLX currently declares backward capture unsupported and raises tiered backend errors. Torch-only features such as autograd node records and live backward intervention remain out of reach unless another backend exposes hookable backward graphs. JAX exposes trace.autograd.derived_grads, populated by tl.backends.jax.GradOptions through a second pure jax.value_and_grad run over fn(params, *inputs), and can optionally expose exact saved-op records through trace.autograd.intermediate_derived_grads plus read-only op.autograd.derived_grad with GradOptions(intermediate_grads=True, max_intermediate_grads=...). The JAX intermediate producer runs as a separate zero-tap AD replay and only exposes oracle-confirmed status == "exact" records. MLX exposes trace.autograd.derived_grads, populated by tl.backends.mlx.GradOptions through a second pure mx.value_and_grad run that rebinds mlx.nn.Module params with model.update(params) and refuses records unless the AD-rerun raw output matches the captured output. MLX can optionally expose exact saved-op records through trace.autograd.intermediate_derived_grads plus read-only op.autograd.derived_grad with GradOptions(intermediate_grads=True, max_intermediate_grads=...). The MLX intermediate producer uses custom-VJP identity taps in that auxiliary AD replay, attaches by unambiguous grouped signatures, skips duplicates, and only exposes oracle-confirmed status == "exact" records. tinygrad exposes trace.autograd.derived_grads, populated by tl.backends.tinygrad.GradOptions through a bracketed DEV=PYTHON leaf-gradient run, and can optionally expose exact unambiguous per-op records through trace.autograd.intermediate_derived_grads plus read-only op.autograd.derived_grad. These are not Paddle exposes trace.autograd.derived_grads, populated by tl.backends.paddle.GradOptions through a second guarded Paddle AD pass over module params and selected inputs, and can optionally expose exact saved-op records through trace.autograd.intermediate_derived_grads plus read-only op.autograd.derived_grad. Paddle traces keep has_backward_pass=False; this is not true backward capture. TensorFlow currently defers true backward capture and T1/intermediate derived gradients; TF traces are forward eager-capture records only. trace.backward(...), trace.autograd.backward_passes, trace.saved_grad_ops, and op.autograd.grads raise on JAX, MLX, tinygrad, Paddle, and TensorFlow traces.

Real per-fire timing is measured on every hooked node (universal path): a lightweight timing prehook and the hook entry stamp pair one time.perf_counter() span per fire, served LIVE by trace.autograd.grad_fn_fire_timings (experimental name; keys match trace.autograd.grad_fn_calls, untimed fires read None, and loaded traces refuse typed with grad_fn_fire_timing_unavailable because the runtime event stream never persists). The persisted GradFnCall timing fields keep their shipped single wall-stamp values until the coordinated tlspec bump activates the paired monotonic semantics together with the grad_fn_timing_provenance discriminator; backward_duration's nullable contract change rides that same bump. Implicit passes journal their close at the engine-drain boundary when torch's final-callback queue is available (the queued callback identity-checks its captured pass index and graph-task id, so a stale callback can never close a newer pass), with the synchronization-point backstop always armed — final callbacks skip the engine's error path, so the drain is opportunistic, never presumed. The close routine is a journal/scavenge/finalize split with the finalize guard IN-ROUTINE: the D2H fence and the full projection never run inside an engine invocation; a deferred finalize runs at the next qualifying sync point or non-engine read, and a read from inside an engine invocation journals without materializing. The BackwardPassEnd sidecar event discloses which close path fired (engine_drain / sync_point, experimental values, runtime-only in wave 2; the projected BackwardPass field was deliberately deferred at the tlspec v8 bump and waits for a future one). The engine's final-callback handle resolves through the compat chokepoint behind the named HAS_AUTOGRAD_ENGINE_QUEUE_CALLBACK capability flag (visible in tl.report.doctor()); without it every implicit close takes the sync-point backstop.