torchscan

This reference covers unreleased development APIs; installation distinguishes it from stable 0.2.0. Start with the runnable quickstart to inspect cost, measure a workload, check a change, and consume the report. Read Model and input support for non-trivial calls.

summary and crawl_module return an AnalysisReport for an evaluation forward. measure_workload and measure_latency return BenchmarkReport evidence while preserving callable state. Timing is clean; FLOPs, scoped memory, and profiler evidence use separate passes. Check statuses, methods, and scopes before comparing or consuming numbers.

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.

Model cost 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.

Workload timing

measure_workload assembles selected FLOPs, timing/throughput, and scoped memory into one report with a readable terminal summary. profile=True adds a separate instrumented pass; RSS requires an explicit rss_command. Use metrics=("latency", "throughput") for timing only and print_summary=False for automation. These APIs require the unreleased development version. See Latency and throughput for work-unit throughput, repeated-call behavior, and measurement boundaries. Configure model state, device placement, precision, and threads yourself. The inputs argument records metadata; it is not forwarded to the workload.

measure_workload

measure_workload(workload: Callable[[], object], *, device: str | device, inputs: Any = None, metrics: Sequence[_WorkloadMetric] = _METRICS, work_units: float = 1, work_unit: str = 'samples', warmup: int = 5, min_run_time: float = 0.2, min_repeats: int = 5, profile: bool = False, rss_command: Sequence[str] | None = None, print_summary: bool = True) -> BenchmarkReport

Collect selected workload evidence and print ms, work units/s, and MiB.

Parameters:

Name	Type	Description	Default
workload	Callable[[], object]	Repeatable zero-argument callable; owns state, gradients, precision, and placement.	required
device	str | device	Device on which the workload already executes.	required
inputs	Any	Optional input metadata, never forwarded to the workload.	None
metrics	Sequence[_WorkloadMetric]	Selected measurements; omitted totals stay unavailable with method not_requested.	_METRICS
work_units	float	Samples (or other units) per call, such as batch size; never inferred.	1
work_unit	str	Throughput unit name, defaulting to samples.	'samples'
warmup	int	Explicit calls before timing, in addition to timer calibration.	5
min_run_time	float	Minimum warmed timing duration in seconds, not a deadline.	0.2
min_repeats	int	Minimum timed blocks, at least two.	5
profile	bool	Collect a separate operator profile after FLOPs and memory.	False
rss_command	Sequence[str] | None	Executable and arguments for fresh-process lifetime RSS; must exit.	None
print_summary	bool	Print a terminal summary in addition to returning the report.	True

Returns:

Type	Description
BenchmarkReport	A JSON-serializable BenchmarkReport accepted by render_report. Partial and unavailable
BenchmarkReport	results retain diagnostics. Workload errors and failed RSS commands propagate.

Notes

Timing runs first; each diagnostic pass invokes the same callable once more. Caller state and random generators are not reset. Keep threads fixed during all passes. First-call time includes all callable work; warmed latency describes block averages, not request percentiles. PyTorch peak memory and whole-process RSS have different scopes. Use the individual collectors for custom FLOP formulas or trace export.

measure_latency

measure_latency(workload: Callable[[], object], *, device: str | device, inputs: Any = None, work_units: float = 1, work_unit: str = 'calls', warmup: int = 5, min_run_time: float = 0.2, min_repeats: int = 5) -> BenchmarkReport

Measure first-call time, warmed block-average latency, and throughput.

Parameters:

Name	Type	Description	Default
workload	Callable[[], object]	Zero-argument callable. Owns model state, inputs, gradient mode, precision, transfers, and any other work included in the timing. It is called repeatedly, including during PyTorch timer calibration.	required
device	str | device	CPU, CUDA, or MPS device on which the workload already executes. Select and place the workload yourself; this argument only chooses synchronization. Work on that device is completed before reading each timer boundary.	required
inputs	Any	Optional caller-supplied input description. Only recursive metadata is recorded; this argument is not forwarded to the workload.	None
work_units	float	Positive number of work units completed by each invocation.	1
work_unit	str	Unit name, such as "images" or "tokens"; throughput uses this unit per second.	'calls'
warmup	int	Explicit warmup calls after the first call. PyTorch also calibrates its timer.	5
min_run_time	float	Positive minimum duration of warmed block measurements in seconds. Calibration, first call, warmup, and minimum repeats can extend total runtime.	0.2
min_repeats	int	Minimum number of timed blocks, at least two.	5

Returns:

Type	Description
BenchmarkReport	Structured timing metrics, block samples, and execution/input metadata.
BenchmarkReport	Latency statistics describe block averages, not individual request percentiles.
BenchmarkReport	First-call time is not fresh-process startup or model-load time.

Raises:

Type	Description
TypeError	If workload is not callable.
ValueError	If timing settings or work units are invalid.
ImportError	If the installed PyTorch benchmark utilities lack their dependencies.
RuntimeError	If the device is unavailable or timing produces invalid samples.
NotImplementedError	If the device backend is unsupported.

Notes

Exceptions and workload side effects are preserved. This function does not change gradient mode, move inputs, evaluate the model, or measure FLOPs/memory. Configure PyTorch threads before building the model. The active count is recorded; keep it unchanged during the workload. Use one device per workload; multi-device/distributed synchronization is not provided.

BenchmarkReport

Bases: TypedDict

JSON-serializable timing report for one caller-controlled workload.

See checked performance experiments for output checks and offline benchmark reports.

compare_benchmarks

compare_benchmarks(before: BenchmarkReport, after: BenchmarkReport, *, check: Callable[[], object]) -> BenchmarkComparison

Compare matching benchmark reports and run an owner-supplied output check.

Parameters:

Name	Type	Description	Default
before	BenchmarkReport	Baseline report from measure_latency, optionally including memory/profile evidence.	required
after	BenchmarkReport	Candidate report for matching hardware, software, inputs, and work units.	required
check	Callable[[], object]	Caller-owned output check, executed outside timing. Exactly None or Python True means success; Python False or AssertionError means failure. Other exceptions propagate.	required

Returns:

Type	Description
BenchmarkComparison	Measured deltas only when the output check passes. Input/hardware/software
BenchmarkComparison	mismatch raises before checking. Changed thread/warmup settings remain visible.
BenchmarkComparison	An IQR-based change label is descriptive, not a significance or accuracy test.

Raises:

Type	Description
ValueError	If reports, contexts, methods, or check results are incompatible.
TypeError	If check is not callable.

BenchmarkComparison

Bases: TypedDict

Checked before/after evidence, without storing workload outputs.

Workload diagnostics

See workload diagnostics for separate process RAM and profiler passes.

profile_workload

profile_workload(workload: Callable[[], object], *, device: str | device, trace_path: str | Path | None = None, limit: int = 20) -> ProfileReport

Find expensive profiler events without treating their timings as benchmark latency.

Parameters:

Name	Type	Description	Default
workload	Callable[[], object]	Zero-argument callable, invoked once with its state unchanged by TorchScan.	required
device	str | device	Device on which the workload already executes. CPU/CUDA are traced; MPS returns CPU dispatch events with an explicit GPU-time limitation.	required
trace_path	str | Path | None	Optional Chrome trace destination. Traces can contain input shapes and execution details; keep them with the experiment that owns the data.	None
limit	int	Positive maximum number of grouped operator rows, ranked by device self time on CUDA and CPU self time otherwise.	20

Returns:

Type	Description
ProfileReport	Grouped operator/runtime events, calls, input shapes, self times in seconds, and net allocated bytes.
ProfileReport	Allocated bytes are event deltas, not peak RAM. Parent/overlapping time must
ProfileReport	not be summed into a claimed model latency.

Raises:

Type	Description
TypeError	If workload is not callable.
ValueError	If limit is invalid.

Notes

Profiling is a separate instrumented pass. Use measure_latency for clean timing. No model evaluation, warmup, device movement, or output retention is added.

ProfileReport

Bases: TypedDict

Operator diagnostics from a separate, instrumented workload call.

Consume saved reports

Serialize returned mappings with json.dumps. render_report renders model or workload HTML without remeasurement; model reports also support SVG. A saved comparison displays its recorded output check without rerunning it.

render_report

render_report(report: AnalysisReport | BenchmarkReport | BenchmarkComparison, *, 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 | BenchmarkReport | BenchmarkComparison	Schema-v1 model analysis, benchmark, or checked benchmark comparison.	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.