Skip to content

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) -> 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.

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) -> AnalysisReport

Print and return a truthful module analysis report.

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 and layer-call metrics 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.

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.

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.