Node labels and encoding channels¶

Every name on this page is experimental (keyword-only) and may change in any release.

All images under docs/images/encoding/ are regenerated by python scripts/render_encoding_reference.py.

Encoding channels¶

One declarative surface: a VALUE SOURCE (record field name, builtin token, or callable node -> value) mapped to a VISUAL CHANNEL, each with a normalization, a scale transform and a plain drawn key, with the details kept on the encoding state (see "The key" below). Channels are strictly presentation-only (the Trace and collapse plan are untouched), strictly opt-in (plain draw() is byte-identical), and dot-layout-only (layout="auto" forces dot when a channel is active; explicit layout="rank" refuses encoding_requires_dot_layout).

color_by — sequential fill¶

draw(color_by="time") fills eligible op nodes from a colorblind-safe, perceptually uniform sequential ramp, normalized linear min-max over visible nodes; the key draws a colour bar labelled with the quantity and its units at both ends. Missing or rolled-ambiguous values leave nodes unencoded with a note on the encoding state.

color_by="device_time" / size_by="device_time" (torchnative.3) read the session-time Kineto join: exact-attribution device nanoseconds per pass-qualified op, summed, with the mandatory aggregation line in the encoding state, on rolled nodes. The channel exists only on a trace captured under an owned profiler session whose join reached joined (torchlens.report.kernel_profile); explicitly requesting it anywhere else refuses typed (device_time_unavailable) with the cause and remedy — never silent zeros. Joined-but-kernel-less nodes read honest n/a.

size_by + scale — node area¶

size_by accepts the same builtin tokens as color_by with the same meaning and units ("time", "bytes", "flops", "magnitude", "grad_norm", "device_time"), plus record field names, callables, and the shape token "dims": draw(size_by="dims", size_scale="sqrt") sizes nodes by the numel of the non-batch output shape (the conservative area-only mapping). On rolled multi-pass nodes a source that cannot be certified single-valued refuses (size_by_rolled_varying).

Sizes are label-aware. Once every label is resolved, the draw estimates each label box and takes the median box among the size-encoded nodes as the base: the smallest-value node is the size of a typical label, and area grows linearly to 4x the base at the largest value. Nodes with typical labels share the base, so their areas compare directly. A node whose label is wider or taller than the median starts from its own label box instead and grows by the same factor, so one long argument row never inflates every node. Sizes are emitted as width/height MINIMUMS under fixedsize=false: labels never truncate and fonts never scale. A node_spec_fn that sets its own width/height keeps them.

Collapsed module boxes and segments¶

When the graph is collapsed (depth=N, collapse="auto", collapse="max"), module boxes and segment capsules are encoded too. A box's value is the EXACT sum of its member ops for additive sources: "time" (func_duration), "bytes" (activation_memory), "flops" (flops_forward), "device_time", and the additive fields transformed_activation_memory, flops_backward, flops_total and macs_forward. The encoding state records the aggregation (module boxes: sum over member ops, ... across all calls on rolled views, segments: sum over member ops). An op with no function (inputs, outputs, buffers) contributes zero time and flops. A box with any member op missing a value is left unencoded with a note on the encoding state rather than shown with a partial sum. Sources with no exact sum over ops ("magnitude", "grad_norm", "dims", other fields, callables) leave every box unencoded, and the encoding state says so (collapsed nodes unencoded: <source> has no exact sum over ops). A repeat fold box shows its representative module only, and the encoding state notes it.

The key¶

The drawn key carries no explanation lines. With color_by it is a small colour bar labelled with the quantity and its units at both ends (time: 0.098 ms [bar] 1.46 ms); with size_by, one line (size: output bytes, 15 KB ... 60 KB); with stack_by, one line (aligned: same <annotation>). When no value range exists the colour key is one line instead (time: all 0.1 ms, time: no values).

The drawing carries nothing else: no transform, coverage or scale wording, drawn or in the DOT source. show_legend=False drops the key.

The ramp is a stop list blended in OKLab, with each stop placed at its share of the path's perceptual length, so equal value steps read as equal colour steps. The default is viridis: on a light page low is yellow and high is dark purple, on the dark page the list runs the other way (high is bright). ColorBrewer YlOrBr (light page: its 9-class list without the near-white first class) is built in as a named alternative, selected by one internal constant (torchlens/_visualization/_ramps.py).

stack_by — rank columns (the classic timestep diagram)¶

draw(stack_by=True, direction="leftright") pins nodes sharing an annotation value to one Graphviz rank. The auto form derives pass_number on multi-pass ops under the LOCKSTEP LICENSE (globally monotone execution windows); the caption discloses the annotation used, and the key line reads aligned: same <annotation>.

A non-lockstep trace (e.g. two sequential loops) REFUSES the auto form (stack_by_auto_underivable) — no ground truth for "same column = same timepoint" exists. An explicit annotation bypasses the license (the user asserts their own semantics, disclosed in the caption):

Across-pass shape summary row¶

A rolled multi-pass layer whose passes vary in output shape carries Layer.shape_summary (L1's derived field: "4->2" monotone, "4-2" min-max, full first->last shapes for multi-axis variation). When set, the default node label renders it VERBATIM as one plain-text row directly after the title row, escaped by the standard choke point. Single-pass layers and unrolled per-pass nodes carry no summary row. The node_label_fields selector token "shape_summary" adds the row to explicit field pickers (skipped when the field is unset).

The "dtype" selector token adds the node's output dtype as one row, framework namespace dropped ("float16"), for reading where a mixed-precision graph changes dtype: trace.draw(node_label_fields=["name", "shape", "dtype"]). Nodes without a dtype contribute no row.

Saved-for-backward annotation (opt-in)¶

draw(show_saved_for_backward=True) makes autograd's memory behaviour visible per node: every op whose grad_fn measurably retained tensors for the backward pass gains one plain-text label row, saved for backward: N tensors, X MB, from the capture-time num_autograd_tensors / autograd_memory measurements (deduped by storage within each op). The row appears ONLY where retention was measured as positive — an op that saved nothing, or whose backward graph was never built (e.g. a torch.no_grad() capture), gets no row, so an absent row makes no claim. On rolled multi-pass layers the stored measurements are already cross-pass sums and the row discloses (total across passes). Strictly opt-in: plain draw() labels are unchanged. Explicit node_label_fields pickers keep full control of their rows (no annotation is injected). Composes with color_by="autograd_memory" to add a ramp over the same quantity.

Checked suppression of redundant constructor args (default-on)¶

Node labels omit a module constructor arg exactly when the check licenses it: the arg value provably equals the captured shape dimension it claims to duplicate on this trace. in_features=32, out_features=10 on a Linear whose input/output rows already show (1, 32) -> (1, 10) says nothing the label does not already say — so it is suppressed. A mismatch or an unavailable shape keeps the arg visible: the mismatch case is exactly the interesting one, so the rule can only reveal more, never hide a discrepancy.

• Eligibility is a closed table of torch nn module families (linear in/out_features, conv in/out_channels, batch/instance-norm num_features, layernorm normalized_shape, embedding embedding_dim, MHA embed_dim). Axis semantics are the PyTorch module contract (channel = logical axis 1; channels_last is a memory format, not a logical shape change).
• Never candidates (not recoverable from I/O shapes): kernel_size, stride, padding, dilation, groups, num_embeddings, num_heads, dropout p.
• Input-side rows check the single incoming data edge's source shape; multi-input or shape-less parents stay visible. Non-torch backends and unrecognized modules are never candidates. Ambiguity never suppresses.
• A rolled multi-pass aggregate whose passes disagree on shape has no single output shape, so its args stay visible while the unrolled per-pass nodes suppress — a deliberate, honest cross-mode difference.
• Detached/standalone records render all args (no trace to check against).

draw(show_redundant_args=True) shows every captured arg.

Before/after (conv):

Before/after (transformer block):

Before/after (attention style, node_spec_fn=torchlens.experimental.node_styles.attention_node_mode):