torchscan

This reference follows the development version on main. Read Model and input support before using non-trivial calls and Understanding results before comparing measurements.

Model analysis

crawl_module

crawl_module(module: Module, input_shape: list[tuple[int, ...]] | tuple[int, ...] | None = None, dtype: dtype | Iterable[dtype] | None = None, *, args: tuple[Any, ...] | None = None, kwargs: Mapping[str, Any] | None = None, device: str | device | None = None, strict: bool = False, mode: Literal['full', 'structure'] = 'full', custom_modules: Mapping[type[Module], ModuleHandler] | None = None, custom_mapping: Mapping[Any, Callable[..., int | float]] | None = None) -> AnalysisReport

Collect a truthful, machine-readable report from one inference forward pass.

Calls sharing a module instance must be serialized because analysis temporarily changes its training state and installs forward hooks.

mode="structure" collects shapes, calls, parameters, and buffers without module formulas or operator dispatch. Unrequested compute totals are unavailable with method not_requested; strict checks only requested metrics.

custom_modules maps module types to scoped ModuleHandler callbacks. The closest class in the module's MRO wins; caller handlers precede built-ins. Declared subtree_metrics are inclusive and suppress descendant estimates for those fields, even when a callback fails. Other fields are module-local. custom_mapping supplies separate operator FLOP overrides to measure_flops. Neither mapping changes global registries; structure mode executes neither.

summary

summary(module: Module, input_shape: list[tuple[int, ...]] | tuple[int, ...] | None = None, wrap_mode: str = 'mid', max_depth: int | None = None, receptive_field: bool = False, effective_rf_stats: bool = False, *, dtype: dtype | Iterable[dtype] | None = None, args: tuple[Any, ...] | None = None, kwargs: Mapping[str, Any] | None = None, device: str | device | None = None, strict: bool = False, mode: Literal['full', 'structure'] = 'full', custom_modules: Mapping[type[Module], ModuleHandler] | None = None, custom_mapping: Mapping[Any, Callable[..., int | float]] | None = None) -> AnalysisReport

Print and return a module report; use mode="structure" for shapes and counts only.

Custom module extensions

See the extension tutorial for complete examples and subtree ownership.

ModuleCall dataclass

Actual forward-call objects supplied to an estimation callback.

Parameters:

Name	Type	Description	Default
module	Module	Module whose forward call just completed.	required
args	tuple[Any, ...]	Complete positional arguments, including containers and non-tensors.	required
kwargs	Mapping[str, Any]	Complete actual keyword arguments; omitted defaults are not inserted.	required
output	Any	Complete output, including containers and non-tensors.	required

These objects are borrowed, not copied. Callbacks must treat them as read-only. The callback runs immediately after forward in evaluation/no-grad mode, with operator dispatch counting suspended. Retaining tensors extends their lifetime.

ModuleEstimates

Bases: TypedDict

Independent module estimates returned by a callback.

Numbers mean complete estimates; None explicitly means unavailable. Structured results preserve partial lower bounds and unavailable states. Omitted fields fall back to built-in leaf formulas or, for composites, child estimates. An omitted field owned by a subtree handler remains unavailable.

ModuleHandler dataclass

Scoped callback and explicit ownership of inclusive subtree estimates.

Parameters:

Name	Type	Description	Default
estimate	Callable[[ModuleCall], ModuleEstimates]	Callback accepting a complete ModuleCall and returning independent ModuleEstimates. Exceptions become report diagnostics.	required
subtree_metrics	frozenset[str]	Names of metrics for which this module owns all work in its executed subtree. Children still produce structural call records, but their owned metric fields are omitted to avoid double-counting. Ownership persists when the callback fails or omits an owned estimate.	frozenset()

Other estimates describe only this module's own work, excluding its children. Registration is supplied to one analysis through custom_modules. Matching uses the closest class in the concrete module's MRO; caller handlers are checked before built-in handlers, even when the caller registers a base class.

Operator FLOPs

measure_flops

measure_flops(workload: Callable[[], Any], *, modules: Module | list[Module] | None = None, custom_mapping: Mapping[Any, Callable[..., int | float]] | None = None) -> FlopReport

Measure workload FLOPs with PyTorch's native operator counter.

Parameters:

Name	Type	Description	Default
workload	Callable[[], Any]	Zero-argument callable invoked exactly once inside the counter.	required
modules	Module | list[Module] | None	Optional module or modules used for hierarchical counts on older supported PyTorch releases.	None
custom_mapping	Mapping[Any, Callable[..., int | float]] | None	Per-call PyTorch operator-to-FLOP formula overrides.	None

Returns:

Type	Description
FlopReport	A versioned report with known counts and diagnostics for every observed uncounted operator.

Raises:

Type	Description
NotImplementedError	If the installed PyTorch counter cannot expose the mapping needed to find missing formulas.
Exception	Any exception raised by workload is propagated unchanged.

Report comparison

compare_reports

compare_reports(before: AnalysisReport, after: AnalysisReport) -> ReportDiff

Compare totals, layer-call metrics and optional token dependencies from two reports.

Layers are matched by their full path and call index. Numeric deltas are only produced when both metric results are complete; incomplete states propagate. Token dependencies are copied as structured before/after evidence, without numeric deltas. Older schema-v1 reports may omit this optional information.

Parameters:

Name	Type	Description	Default
before	AnalysisReport	Earlier analysis report.	required
after	AnalysisReport	Later analysis report.	required

Returns:

Type	Description
ReportDiff	A deterministic, JSON-serializable report difference.

Raises:

Type	Description
ValueError	If schema versions or metric methods differ, or a layer-call identity is duplicated.

Offline visual reports

render_report

render_report(report: AnalysisReport, *, format: Literal['html', 'svg'] = 'html', before: AnalysisReport | None = None, title: str = 'TorchScan analysis', metric: View = 'module_flops') -> str

Render structured analysis as a self-contained HTML report or standalone SVG.

Parameters:

Name	Type	Description	Default
report	AnalysisReport	Schema-v1 report returned by crawl_module or summary.	required
format	Literal['html', 'svg']	"html" for interactive native controls or "svg" for a static snapshot.	'html'
before	AnalysisReport | None	Optional baseline passed to compare_reports(before, report).	None
title	str	Plain-text report title. Untrusted strings are escaped.	'TorchScan analysis'
metric	View	Initial HTML view or SVG cost map: module_flops, macs, dmas, parameters, or parameter_bytes.	'module_flops'

Returns:

Type	Description
str	UTF-8-compatible document text. Save with Path("report.html").write_text(result, encoding="utf-8");
str	open with webbrowser.open(Path("report.html").resolve().as_uri()). No server or extra dependencies are needed.

Raises:

Type	Description
ValueError	For invalid schema, metric, format, non-JSON data, invalid measurement states, duplicate call identities, or incompatible comparison methods, units or scopes.

Notes

Inputs are never mutated or remeasured. Execution deltas are withheld when input metadata, execution mode, device/dtype metadata or software versions differ or are missing. Storage totals need compatible complete metrics but do not require identical execution inputs. The SVG exports the module map and selected evidence; HTML adds keyboard selection, zoom, collapse and views.

Public report types and errors

AnalysisReport

Bases: TypedDict

Versioned, JSON-serializable result of a module analysis.

LayerReport

Bases: TypedDict

Information collected for one module invocation.

MetricResult

Bases: TypedDict

JSON-serializable result for one measurement.

Diagnostic

Bases: TypedDict

Machine-readable explanation of an incomplete measurement.

FlopReport

Bases: TypedDict

JSON-serializable operator FLOP report.

ReportDiff

Bases: TypedDict

JSON-serializable differences between two compatible analysis reports.

IncompleteAnalysisError

Bases: RuntimeError

Raised when strict analysis encounters incomplete metrics.