Skip to content

Smart-collapse visual reference

The images on this page are generated by scripts/render_collapse_reference.py. Run PYTHONPATH=$PWD python scripts/render_collapse_reference.py after rendering changes; it updates this gallery and stages the same files in /tmp/collapse_renders/ for review.

The mode ladder

The same ResNet-18 execution is shown at each string mode. none preserves the full graph, auto chooses the first schedule point in the readable band, and max applies the strongest legal condensation.

auto then evens that point out, and only ever condenses further. A small module left expanded beside boxed siblings is boxed too, when it hides no more ops than the smallest box beside it (convnext's stem and stage-2 downsample), so stray raw ops do not sit between boxes. "Beside" means a sibling box or a box one level inside a sibling; a box nested deeper (BERT's encoder.layer.0 next to embeddings) leaves its neighbours expanded. A run of three or more structurally identical, consecutive sibling boxes folds into one segment (convnext's nine stage-2 blocks become stages.2.blocks.0-8 -- 9 blocks). The evened plan never shows more nodes than the schedule point it started from; the ResNet-18, GPT-2 and BERT auto overviews are unchanged.

Full ResNet-18 graph

collapse="none": every operation is visible, including the repeated residual stages that make an overview hard to scan.

Readable ResNet-18 overview

collapse="auto": representative module boxes preserve the stem-to-stage-to-head story while their labels account for hidden work.

Maximum ResNet-18 condensation

collapse="max": the same trace is condensed as far as the renderer can do honestly; it is not a different execution.

The float-t schedule

A float t in [0.0, 1.0] selects the deterministic monotone collapse schedule. 0.0 equals "none", 1.0 equals "max", and increasing t never expands a previously hidden address.

Schedule at t 0

t=0.0: the full graph endpoint.

Schedule at t 0.25

t=0.25: the first intermediate schedule point.

Schedule at t 0.5

t=0.5: a further-condensed view of exactly the same execution.

Schedule at t 0.75

t=0.75: hidden module addresses remain a superset of those at lower t.

Schedule at t 1

t=1.0: the maximum-collapse endpoint.

fold_repeats

Run folding is independent of module collapse. These two views use the same eight-block residual stack with collapse="none"; the right-hand view folds only the repeated sibling run.

Unfolded repeated residual blocks

fold_repeats=False: all eight distinct residual-block instances are shown.

Folded repeated residual blocks

fold_repeats=True: one representative remains and the ellipsis accounts for the other same-class siblings; the stem and head remain at full detail.

Segment boxes

collapse="max" can replace a legal consecutive sibling interval with a dashed segment box. Legal means the members ran back to back in execution order and each feeds the next directly: a box never groups same-kind siblings that alternate with other siblings, and never hides a parent-owned op between its members, so the collapsed graph shows no edge or cycle the traced network lacks.

Maximum collapse on ResNet-50

The dashed segment boxes in this ResNet-50 maximum-collapse view name an address range and totals. They assert adjacency only: a segment is not a real module and never pretends mixed content has one module class.

Output bundles

A decoder that returns its cache returns one output per cached tensor, and each would otherwise draw as its own node. In an overview render (collapse="auto", "max" or a level above 0, unrolled, without show_containers), three or more model outputs that share a container key and named role, dtype and shape draw as one node, such as OUTPUTS: past_key_values.keys with 12 tensors, 1x2x8x8 each, when one top-level module made all of them and that module is drawn as one box. When the module is expanded, its outputs come from several drawn nodes and every output is drawn on its own again. A named single output such as logits or last_hidden_state is never bundled, and collapse="none" never bundles. The bundle counts as one node in the readable band.

What (xN) and +N more promise

These labels look similar but make different, deliberately strict promises about hidden content.

Rolled weight-tied recurrence

(xN) is a true recurrence multiplier: the same parameters were applied N times. This rolled RNN-cell loop has weight reuse, not N independent blocks.

Folded distinct residual blocks

+N more Class means N additional distinct same-class instances, each with its own parameters. The representative box describes one instance; the ellipsis makes the omitted sibling count explicit. Collapsed-box remainders and segment totals likewise count everything they hide, including buffer leaves. This is the honesty contract: labels never claim a repeated call, module, or total that the hidden content does not support.

The compute ceiling

Smart-collapse admission is now gated on the RENDERED UNIVERSE, not raw op count (collapse): the gate variable is U — the full plan's rendered node count after module= focus, depth=, and rolled reductions — tiered by a measured (U, W) work estimator (0.0095 * U**1.43 * W**1.02 ms, W = widest rendered sibling group). COLLAPSE_OPTIMIZER_MAX_OPS (2000) survives as the defensive constant: the quality planner is admitted when U <= 2000 AND the predicted cost fits the budget.

  • Over-budget requests degrade to a deterministic compact fallback plan (significance-greedy boxes plus standard fold discovery), disclosed with the coded collapse_budget_fallback warning and OptimizerResult.planner="linear_fallback" — never an uncollapsed wall.
  • Only PATHOLOGICAL inputs (raw ops above 20x the constant) still decline outright: draw warns (collapse_pathological_skip) and renders a bounded top-level module overview plus an expansion manifest (uncollapsed only when the root has no child module), a forward drawing's Drawing.collapse_plan refuses typed (collapse_plan_unavailable, reason collapse_ops_ceiling), and its Drawing.collapse_schedule degrades to a single full-graph step.
  • Over-ceiling residuals (experimental): when the served overview is still above the readable band -- a floor-fallback plan, or the pathological top-level overview -- draw writes <outpath>.expansion.json beside the picture (schema torchlens_expansion_manifest_v1: one entry per collapsed unit with a stable expansion_id and the draw() keywords that open it) and warns collapse_expansion_manifest with its path.
  • A generous wall-clock watchdog (off in CI and under TORCHLENS_DETERMINISTIC; TORCHLENS_COLLAPSE_WATCHDOG=0 disables) may abandon the quality planner mid-run to the same fallback; every firing is disclosed as an estimator bug (collapse_watchdog_fallback, with U, W, predicted, and actual).

Context reductions are therefore real remedies for the first time: focusing with module=, bounding depth=, or rendering rolled genuinely re-admits the quality planner. The gate is a compute guard, not a correctness limit — the trace itself is complete.

The decline is disclosed, never silent. To get a collapsed view of a very large trace, reduce the rendered graph first with module= focus, depth=, or rolled mode; fold_repeats and show_containers styling are unaffected.

Drawing.collapse_plan diagnostics

The Drawing a forward trace.draw(...) returns carries the diagnostics for the picture it rendered. Drawing.collapse_plan is the renderer-faithful plan the drawing was built from: its mode is the draw call's collapse= ("auto", "max", or a float level) and its context is the render request the draw resolved, so rolled or unrolled, module= focus, depth=, and theme all follow the draw call. Drawing.collapse_schedule exposes the ordered float steps, targets, visible counts, and collapsed module addresses under that same context. Use these diagnostics to inspect a rendering decision rather than inferring it from pixels.

import torch
from torch import nn

import torchlens as tl

model = nn.Sequential(nn.Linear(4, 4), nn.ReLU()).eval()
trace = tl.trace(model, torch.randn(2, 4))
drawing = trace.draw(collapse="auto", file_format="dot")
plan = drawing.collapse_plan
schedule = drawing.collapse_schedule
level_plan = trace.draw(collapse=0.5, file_format="dot").collapse_plan

Both properties refuse typed (collapse_plan_unavailable) for a backward or combined graph, an animation, or once the source trace has been released; collapse_plan also refuses for a drawing made with collapse="none", which has no plan. To compare modes, make one draw call per mode.

Collapse plan and schedule diagnostics

This image is generated from the same ResNet-18 trace as the ladder. It shows real auto and max plans followed by every schedule step; visible counts are monotone as t increases.