Skip to content

fix(summary): support sparse tensor inputs in memory accounting - #418

Merged
TylerYep merged 2 commits into
TylerYep:mainfrom
DangTAnh:pr/150-sparse-input-memory
Oct 11, 2026
Merged

TylerYep merged 2 commits into
TylerYep:mainfrom
DangTAnh:pr/150-sparse-input-memory

Conversation

@DangTAnh

@DangTAnh DangTAnh commented Oct 2, 2026

Copy link
Copy Markdown
Contributor

Summary

Models that take sparse tensors (e.g. GCN adjacency matrices) crashed summary() with NotImplementedError: Cannot access storage of SparseTensorImpl, so no summary — including the Param # column from #150 — was produced. Sparse layouts expose no single storage, so input-memory accounting now sums the member tensors (indices + values) instead.

Changes

  • torchinfo/torchinfo.py: new get_tensor_memory helper with sparse fallback; get_total_memory_used uses it.
  • tests/torchinfo_test.py: new test_sparse_input regression test (GCN with sparse COO adj).

Testing

  • pytest tests/torchinfo_test.py tests/exceptions_test.py tests/mapping_test.py --no-output — 94 passed.
  • Repro from Summary doesn't give the number of parameters for a GCN #150 (GCN + sparse adj): crashed before, now prints Param # correctly (172/172 params); verified COO, CSR, CSC, and dict inputs.
  • 3 golden-file errors in the full run are pre-existing torch-version drift (fail identically on clean tree).

Related

Fixes #150

@DangTAnh
DangTAnh force-pushed the pr/150-sparse-input-memory branch 2 times, most recently from 16083b9 to 17a5992 Compare October 9, 2026 02:00
Sparse COO/CSR/CSC adjacency inputs (e.g. GCNs) crashed get_total_memory_used via untyped_storage, so no summary was produced. Sum member-tensor storage for sparse layouts instead.

Fixes TylerYep#150
@DangTAnh
DangTAnh force-pushed the pr/150-sparse-input-memory branch from 6eaccaf to c171f37 Compare October 9, 2026 02:01
@TylerYep

Copy link
Copy Markdown
Owner

Thanks!

@TylerYep
TylerYep merged commit d67f0d0 into TylerYep:main Oct 11, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Summary doesn't give the number of parameters for a GCN

2 participants