torchscan 0.2.0

TorchScan inspects a PyTorch model and returns a JSON-serializable report of its structure, parameters, inputs, module estimates, and operator FLOPs. Every metric says whether it is complete, partial, or unavailable, so an unsupported operation cannot masquerade as zero.

Quickstart

import torch.nn as nn
from torchscan import crawl_module, summary

model = nn.Conv2d(3, 8, 3)

# Print the human-readable table and receive the same structured report.
report = summary(model, (3, 32, 32))

# Or collect the report without printing the table.
report = crawl_module(model, (3, 32, 32), strict=True)

summary keeps the familiar terminal UX while returning the structured report:

__________________________________________________________
Layer     Type      Output Shape      Param #    Trainable
==========================================================
conv2d    Conv2d    (1, 8, 30, 30)    224        True
==========================================================
Trainable params: 224
Non-trainable params: 0
Total params: 224
----------------------------------------------------------
Model size (params + buffers): 0.00 Mb
----------------------------------------------------------
Module-formula forward FLOPs: 388.80 kFLOPs
Multiply-Accumulations: 194.40 kMACs
Direct memory accesses: 201.82 kDMAs
Operator forward FLOPs: 388.80 kFLOPs
__________________________________________________________

input_shape excludes the batch dimension. For realistic calls—including masks, scalars, None, and nested containers—pass complete args and kwargs instead:

import json

import torch
from torch import nn
from torchscan import crawl_module


class MaskedModel(nn.Module):
    def forward(self, input_ids, *, attention_mask):
        return input_ids * attention_mask


transformer_model = MaskedModel()
input_ids = torch.ones(1, 4)
attention_mask = torch.tensor([[True, True, False, False]])
report = crawl_module(
    transformer_model,
    args=(input_ids,),
    kwargs={"attention_mask": attention_mask},
)
print(json.dumps(report["inputs"]["kwargs"]["attention_mask"], indent=2))

Only metadata is retained:

{
  "kind": "tensor",
  "shape": [1, 4],
  "dtype": "torch.bool",
  "device": "cpu",
  "requires_grad": false
}

TorchScan temporarily evaluates the model with gradients disabled and restores every module's original training state. It records input metadata, never tensor values.

For shapes and parameter counts, skip compute analysis with one option:

report = summary(model, (3, 32, 32), mode="structure")

Structure mode collects the same hierarchy, calls, input/output metadata, parameters, and buffers without FLOP dispatch or module formulas. Unrequested compute totals have status="unavailable" and method="not_requested". strict=True checks the requested metrics. Full analysis remains the default. Both modes release intermediate activations as execution progresses.

Workload measurements

Use zero-argument callables when the owner needs full control over execution:

import json

import torch
from torchscan import measure_flops
from torchscan.process import measure_peak_memory

inputs = torch.ones(8)
flops = measure_flops(lambda: torch.sin(inputs))
print(json.dumps(flops["total"], indent=2))
print("uncounted operator:", flops["diagnostics"][0]["operator"])

memory = measure_peak_memory(lambda: torch.cos(inputs), device=inputs.device)
print(memory["device"], memory["metric"])

measure_flops uses PyTorch's operator dispatch. measure_peak_memory invokes the workload exactly once and reports backend-specific PyTorch memory—not process RSS or total device memory.

Here, PyTorch has no built-in aten.sin formula, so TorchScan shows a lower bound instead of a false zero:

{
  "status": "partial",
  "value": null,
  "known_value": 0,
  "unit": "FLOPs",
  "scope": "workload",
  "method": "torch.utils.flop_counter.FlopCounterMode"
}
uncounted operator: aten.sin
cpu pytorch_tensor_bytes

Peak byte values are intentionally omitted because they depend on the workload, allocator, PyTorch version, and hardware; the returned mapping includes baseline_bytes, peak_bytes, and delta_bytes.

Before/after comparison

import torch.nn as nn
from torchscan import compare_reports, crawl_module

before = crawl_module(nn.Conv2d(3, 8, 3), (3, 32, 32))
after = crawl_module(nn.Conv2d(3, 12, 3), (3, 32, 32))
diff = compare_reports(before, after)
parameters = diff["totals"]["parameters"]
print(parameters["status"], parameters["delta"])

complete 112

compare_reports propagates incomplete metrics. It does not store baselines or decide whether a model fits a budget; the model owner supplies those policies.

Trust the status, not only the number

• complete: the requested scope was counted; value is authoritative for the documented method.
• partial: known_value is a lower bound and diagnostics identify missing work.
• unavailable: TorchScan cannot produce the metric for this execution.

Use strict=True when any incomplete analysis must stop automation. See the report schema and methodology before comparing results.

Installation

The examples here use the development API on main, which requires Python 3.11+ and PyTorch ≥2.1,<3:

pip install git+https://github.com/frgfm/torch-scan.git

For the published stable release, use pip install torchscan. Its API and requirements differ; see the installation guide and v0.2 migration guide. For a local development checkout, follow Contributing.

Documentation

• Agent quickstart
• Model and input support
• v0.2 migration guide
• API reference

Agents can also load the repository skill at .agents/skills/torchscan/SKILL.md.

Citation

Citation metadata is available in CITATION.cff.

Contributing and license

Contributions are welcome; see CONTRIBUTING.md. TorchScan is distributed under the Apache License 2.0.