tl.export reference¶
tl.export writes a completed trace to static graph files, performance-viewer formats,
tabular files, tracker objects, and an xarray activation assembly. File exporters return the
written path; tracker exporters accept existing tracker objects and do not create runs. The
examples use DOCS_TMPDIR when the documentation test provides it, otherwise /tmp.
Static graph files¶
svg¶
tl.export.svg(log, path, *, editable=True) writes a lightweight SVG graph. editable=True
adds stable IDs and semantic CSS classes. html provides a self-contained interactive
viewer of the same graph data.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.svg(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "trace.svg")
print(path.name, path.exists())
Output:
html¶
tl.export.html(log, path) writes a self-contained HTML viewer with pan, zoom, and node-hover
support; it has no network dependency. See svg for the lighter static output.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.html(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "trace.html")
print(path.name, path.exists())
Output:
model_explorer¶
tl.export.model_explorer(log, path) writes a static JSON graph collection with nodes, edges,
labels, shapes, and memory fields that Google Model Explorer opens directly (File > Open —
the top-level label/graphs shape and node fields are pinned against the
ai-edge-model-explorer 0.1.32 ingest contract in tests/test_exports.py). It is a data
export, not a runnable model.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.model_explorer(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "graph.json")
print(path.name, path.exists())
Output:
netron¶
tl.export.netron(log, path) writes an enriched, lossy ONNX ModelProto JSON graph
(schema v2, irVersion 10) that Netron opens fully lit: typed shapes on every edge, native
graph inputs/outputs, the module hierarchy as drill-down FunctionProtos (the default
granularity="module" view), curated panel facts, a netron:attachment metrics
companion written beside the file by default (attachment=False opts out), and, with
baseline=<reference Trace>, per-node diff facts from torchlens.report.diff.
The acceptance contract — strict protobuf-JSON parse into
onnx.ModelProto plus onnx.checker.check_model(full_check=True) — is pinned in
tests/test_netron_export_contract.py, and Netron 9.2.2's own parser is executed over the
artifact in tests/test_netron_export_vendor.py. It deliberately is not a runnable ONNX
model: ops keep their captured TorchLens names under the custom ai.torchlens.lossy
domain; use it only for graph inspection. Full guide: docs/reference/netron_export.md.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.netron(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "netron.json")
print(path.name, path.exists())
Output:
Performance formats¶
chrome_trace¶
tl.export.chrome_trace(log, path) writes a Chrome tracing JSON timeline for one trace. Open the
file in Chrome's tracing viewer or another Chrome-trace-compatible tool. For a bundle, use
chrome_trace_diff.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.chrome_trace(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "trace.json")
print(path.name, path.exists())
Output:
chrome_trace_diff¶
tl.export.chrome_trace_diff(bundle, path) writes one Chrome timeline comparing the members of a
TorchLens Bundle. Its per-node deltas are the relative L2 distances from the baseline (what bundle.compare(baseline, member, metric="relative_l2") reads per pair), so they need members that pass the
Bundle comparison gate (one shared graph, and sites that line up: equal input values, or for
members on different inputs equal graph-shape hashes) and equal shapes at every compared site.
Members that do not still export their timelines; only the deltas are refused: every event's
delta is None and metadata["deltas"] records the refusal code (for example
bundle_gate_input_values_differ for graphs that do not line up, or metric_shape_mismatch for
inputs of different shapes).
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
x = torch.ones(1, 2)
bundle = tl.Bundle({"first": tl.trace(nn.ReLU(), x), "second": tl.trace(nn.ReLU(), x)})
path = tl.export.chrome_trace_diff(bundle, Path(globals().get("DOCS_TMPDIR", "/tmp")) / "diff.json")
print(path.name, path.exists())
Output:
speedscope¶
tl.export.speedscope(log, path) writes an evented speedscope
profile using TorchLens operation durations.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.speedscope(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "profile.speedscope.json")
print(path.name, path.exists())
Output:
flamegraph¶
tl.export.flamegraph(log, path) writes folded-stack text containing operation durations. Feed
that text to a flamegraph renderer; it is not itself an SVG.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.flamegraph(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "stacks.folded")
print(path.name, path.exists())
Output:
memory_timeline¶
tl.export.memory_timeline(log, path) writes a tensor-scope memory timeline. It reports tensor
bytes observed and retained by TorchLens, not allocator or process peak memory.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.memory_timeline(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "memory.json")
print(path.name, path.exists())
Output:
Tabular and activation formats¶
csv¶
tl.export.csv(log, path, **kwargs) writes Trace.to_pandas() as CSV. Extra keywords are
forwarded to DataFrame.to_csv; json and parquet write the same table
in other forms.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.csv(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "trace.csv")
print(path.name, path.exists())
Output:
parquet¶
tl.export.parquet(log, path, **kwargs) writes Trace.to_pandas() as Parquet. It requires
pyarrow (install torchlens[tabular] if it is unavailable).
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.parquet(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "trace.parquet")
print(path.name, path.exists())
Output:
json¶
tl.export.json(log, path, *, orient="records", **kwargs) writes Trace.to_pandas() as JSON.
orient and extra keywords are forwarded to pandas.
from pathlib import Path
import torch
from torch import nn
import torchlens as tl
path = tl.export.json(tl.trace(nn.ReLU(), torch.ones(1, 2)), Path(globals().get("DOCS_TMPDIR", "/tmp")) / "trace.json")
print(path.name, path.exists())
Output:
xarray¶
tl.export.xarray(log) returns an xarray.Dataset. Its out variable is the
NeuroidAssembly-shaped xarray.DataArray of saved non-buffer tensor outputs with
presentation and neuroid dimensions; all of them must have the same presentation count.
Every saved buffer layer (module state such as BatchNorm running statistics, retained whatever
save= selects) is its own data variable named by its layer label, with its own dimensions
<label>_dim_<i>, no presentation axis, its values unchanged, and the variable attribute
torchlens_kind="module_state" (plus layer_label and buffer_address);
dataset.attrs["module_state_variables"] lists them. A save-all trace and a partial trace
export buffers the same way; a partial trace's unsaved layers get no data variable and are
listed in attrs["skipped_unsaved_layers"]. xarray must be installed.
import torch
from torch import nn
import torchlens as tl
dataset = tl.export.xarray(tl.trace(nn.ReLU(), torch.ones(1, 2)))
assembly = dataset["out"]
print(assembly.dims, assembly.shape)
Output:
Tracker integrations¶
tensorboard¶
tl.export.tensorboard(log, writer, step=<int>, prefix="torchlens") writes scalar and text
summaries to an existing object with TensorBoard's add_scalar interface, then returns that
writer. step is required and keyword-only: every emission path carries the caller's own
global step (the historical step=0 default silently piled every export onto one x-axis
coordinate and was removed).
import torch
from torch import nn
import torchlens as tl
class Writer:
def __init__(self): self.calls = []
def add_scalar(self, *args): self.calls.append(args)
def add_text(self, *args): self.calls.append(args)
def flush(self): pass
writer = Writer()
tl.export.tensorboard(tl.trace(nn.ReLU(), torch.ones(1, 2)), writer, step=0)
print(len(writer.calls))
Output:
wandb¶
tl.export.wandb(log, run=None, name="torchlens_trace") creates a Weights & Biases table and
logs it only if an existing run is supplied or active. It requires the wandb extra.
import torch
from torch import nn
import torchlens as tl
result = tl.export.wandb(tl.trace(nn.ReLU(), torch.ones(1, 2)))
print(sorted(result))
Output:
mlflow¶
tl.export.mlflow(log, client=None, prefix="torchlens", *, run_id=None) logs, to an existing
run: the three summary metrics, the per-layer statistics table as the artifact
<prefix>/layers.json, the report (honesty preamble, fact block, trace summary) as the text
artifact <prefix>/report.txt, and every capture-honesty fact as a string run tag
<prefix>.capture_honesty.<fact> (strings verbatim, other facts as compact sorted JSON; a tag
over MLflow's 5000-character limit is replaced by a marker naming the report). Pass the fluent
mlflow module to log into the active run (or run_id= of another run), or an
mlflow.MlflowClient together with run_id=. The client must offer log_metric, log_table,
log_text and set_tag, or the call raises TypeError before writing anything. It returns the
metrics plus table, report, tags, artifacts and capture_honesty, with or without a
client.
import torch
from torch import nn
import torchlens as tl
metrics = tl.export.mlflow(tl.trace(nn.ReLU(), torch.ones(1, 2)))
print(sorted(metrics))
Output:
['artifacts', 'capture_honesty', 'num_layers', 'num_saved_ops', 'report', 'table', 'tags', 'total_activation_memory']
aim¶
tl.export.aim(log, run=None, prefix="torchlens") tracks the three summary metrics on an
existing run, writes the per-layer table as the run param <prefix>.layers (one dict per layer)
and as one series per numeric column (<prefix>.layers.<column>, context {"table": "layers"},
step = row index), the report as the run param <prefix>.report, the full fact block as
<prefix>.capture_honesty and each fact as a run tag <prefix>.capture_honesty.<fact>=<value>.
The run must offer track, add_tag and item assignment. It returns the metrics plus table,
report, tags and capture_honesty.
import torch
from torch import nn
import torchlens as tl
metrics = tl.export.aim(tl.trace(nn.ReLU(), torch.ones(1, 2)))
print(sorted(metrics))
Output: