Notebook cards and the treescope bridge¶
(Every name on this page is experimental: it may change in a later release.)
TorchLens objects carry dependency-free rich HTML cards as their primary
notebook experience; the treescope bridge is an opt-in upgrade that adds
interactive array visualization. This ordering is not a preference but a
mechanical fact: a live IPython kernel calls _repr_html_ before treescope's
machinery ever runs, so the card IS the top-level experience whether or not
treescope is installed. Upstream treescope is idle (last release v0.1.10 on
2025-08-08, last push 2026-06-18, checked 2026-11-03); the bridge keeps working against
that release and the cards do not depend on it.
Notebook cards¶
Print any of these in a notebook cell and the card renders — stdlib-only generation, no IPython or notebook extra required, readable with JavaScript disabled:
| Object | Card |
|---|---|
Trace |
model/backend/counts, save coverage, outcome badge, NaN/Inf disclosure, params/FLOPs, budgeted foldable index of copyable lookup keys |
Layer |
the text card's first line and rows; repeated layers show up to three per-pass value summaries |
Op |
short key, tensor geometry, value words, position, module, output size, graph, function and source location |
Param |
live tensor geometry, trainability, basis and version, size, cached values, consuming sites and gradient facts |
PartialTrace |
failure-first: FAILED CAPTURE banner, phase/reason, escaped exception, committed prefix and frontier, what remains safe to inspect |
Trace and failed-capture cards carry the PROOF row — capture outcome, verification badge, nonfinite coverage basis, save policy. A TorchLens card asserts "this is a verified execution", not "this is an object".
Hard rules: a card never raises (internal failures degrade to a one-line
card unavailable: <reason>); honesty states are never folded; display paths
never materialize disk payloads and never do unbounded device transfers.
Op, Layer and Param notebook cards show the same first line and labelled
rows as str(record), with at most 10 body rows and no automatic tensor
preview. Saved CPU activations under the sampling gate may compute statistics
once, cached by version; live parameters and device tensors only reuse valid
cached statistics. Missing evidence reads values not computed.
Multi-pass Layer cards show up to three individual passes, never a pooled
statistic. A record card stays below 20 KB.
Copy-expression roots¶
Trace index keys are executable access expressions (log['conv2d_1_1']). The root
name resolves through a ladder: an explicit root= wins; nested treescope
renders use treescope's own path; in a live notebook an identity scan of the
user namespace picks the shortest non-underscore binding (IPython history
names _, _N, _i*, In, Out are excluded — without the exclusions the
tie-break picks _, rebound by the next cell); with no root, keys copy as
literals and the card says why. Offline artifacts never scan a namespace.
The treescope bridge¶
import torchlens as tl
from torchlens.bridge import treescope as tl_treescope
tl_treescope.register() # idempotent; typed conflict on occupied slots
tl_treescope.status() # version, handlers, tensor capability, suppression state
tl_treescope.display(op) # the explicit door: interactive arrayviz when capable
with tl_treescope.disabled(): # scope: treescope's own default rendering
...
tl_treescope.unregister() # removes exactly our handlers, by identity
What the bridge buys: without it, an explicit treescope.display(trace) or a
Trace nested in a rendered container reflects the dataclass field-by-field
into a megabyte internals dump (measured 1.3–2.75 MB, up to 82.9 s on
gpt2/128 tokens). With it: bounded foldable summaries with lookup-key indexes
(~10 KB), and — when the installed treescope can render tensors — the REAL
saved payload through treescope's interactive array visualizer with TorchLens
truncation budgets.
Rules the bridge keeps:
- Never
register_as_default(), never atorch.Tensorhandler — the user keeps control of global display state, and treescope owns the tensor slot. - Container rule: bridged containers emit scalars, strings, and a
budgeted key index with
shown K of N; they never recurse into child TorchLens objects (the recursive version measured WORSE than no bridge). - Capability probe, never a version pin: the only treescope release on
PyPI (0.1.10) reads the removed
Tensor.namesattribute and raises on every torch>=2.13 tensor. The bridge probes behaviorally at first use; a broken adapter means strict degradation — the object card stays, the reason is named, and atreescope_bridge_degradedwarning fires once. - Top-level display prefers our card (
_repr_html_short-circuits treescope).tl_treescope.display(obj)is the explicit door to the interactive render. - The duplicate box: treescope appends any object's
_repr_html_as a collapsed "Rich HTML representation" box after our handler renders. The cards suppress the duplicate payload with a narrow, fail-safe check (a treescope postprocessor frame on the stack AND the bridge having just rendered that exact object); if the check ever misfires, the worst case is a redundant collapsed box, never a missing card.
Version compatibility¶
| treescope | torch | behavior |
|---|---|---|
| absent | any | cards only; register()/display() refuse typed (treescope_bridge_unavailable) |
| 0.1.10 (released) | < 2.13 | full bridge incl. tensor leaves (upstream adapter works) |
| 0.1.10 (released) | >= 2.13 | bridge with STRICT DEGRADATION on tensor leaves (Tensor.names crash; probe-detected), object summaries still bounded |
git main / first release with the names fix |
any supported | full bridge incl. tensor leaves |
Remote capture recipe (single-file report)¶
# on the cluster
python -c "
import torch, torchlens as tl
from mymodel import build; model, x = build()
log = tl.trace(model, x)
tl.export.html(log, 'run_report.html') # one file, zero network
"
scp cluster:run_report.html . # open from file://, no Python needed
See docs/guides/offline_report.md for the report's anatomy and options.
Credits¶
The card and report conventions adopt ideas from treescope / penzai
(Daniel D. Johnson, Google DeepMind; Apache-2.0): external type
registries, collapsed/expanded balanced layouts, copy-buttons that emit
code, faceted array visualization conventions, truncation budgets with edge
items, diverging-around-zero defaults, and pattern-coded exceptional values.
Named code borrow: infer_balanced_truncation (including doubling_bonus)
and the two-stage slice-then-convert truncation protocol, ported with license
notice in torchlens/report/_notebook/_truncation.py.