Support sparse tensors by justinchuby · Pull Request #440 · onnx/ir-py

justinchuby · 2026-06-11T19:00:42Z

Summary

Adds first-class support for sparse tensors in the ONNX IR, including a new SparseTensor class (COO format following SparseTensorProto), serde support, sparse initializers, and sparse-tensor attributes.

This continues the work from #426 with a refined API.

Sparse tensor support

New SparseTensor class in _core.py storing values, indices, and dims (COO format), with numpy() (→ scipy.sparse.coo_array) and as_tensor(lazy=) (→ dense Tensor) conversions.
SparseTensor.from_scipy_sparse() classmethod to construct from any scipy.sparse array (converted to COO). ir.tensor() stays dense-only with a clean TensorProtocol return type.
Value.const_sparse_value field for sparse initializers, with correct name propagation and string representation.
Serde: serialize_sparse_tensor/serialize_sparse_tensor_into/deserialize_sparse_tensor, sparse initializer (de)serialization, and SPARSE_TENSOR/SPARSE_TENSORS attribute handling.
SparseTensorProtocol added to _protocols.py.

Design notes

SparseTensor intentionally has no doc_string/metadata_props because SparseTensorProto only stores values/indices/dims; meta is available for in-memory analysis (not serialized).

Public API & docs

Exposes SparseTensor, AttrSparseTensor, AttrSparseTensors, and the new serde functions.
Documents sparse tensor usage in docs/tensors.md.

Testing

Comprehensive tests for construction, conversion, round-tripping, initializers, and attributes.

codecov · 2026-06-11T19:02:25Z

Codecov Report

❌ Patch coverage is 75.86207% with 35 lines in your changes missing coverage. Please review.
✅ Project coverage is 85.47%. Comparing base (af4bdfe) to head (07a11fc).
⚠️ Report is 19 commits behind head on main.

Files with missing lines	Patch %	Lines
src/onnx_ir/_core.py	78.65%	13 Missing and 6 partials ⚠️
src/onnx_ir/serde.py	70.37%	11 Missing and 5 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #440      +/-   ##
==========================================
- Coverage   85.53%   85.47%   -0.06%     
==========================================
  Files          53       53              
  Lines        7029     7166     +137     
  Branches     1459     1483      +24     
==========================================
+ Hits         6012     6125     +113     
- Misses        650      670      +20     
- Partials      367      371       +4

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

- Add SparseTensor class to _core.py implementing SparseTensorProtocol - Update Value.const_value to accept SparseTensorProtocol - Add deserialize_sparse_tensor/serialize_sparse_tensor/serialize_sparse_tensor_into to serde.py - Update _deserialize_attribute for SPARSE_TENSOR/SPARSE_TENSORS types - Update _fill_in_value_for_attribute for SPARSE_TENSOR/SPARSE_TENSORS types - Update _deserialize_graph to handle sparse_initializer - Update serialize_graph_into to serialize sparse initializers - Update from_proto and to_proto to handle SparseTensorProto - Export SparseTensor and new serde functions in __init__.py - Update docs/tensors.md to document sparse tensor support - Add comprehensive tests for sparse tensor functionality Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

… isinstance checks - Add `name: str | None` to SparseTensorProtocol in _protocols.py - Replace hasattr checks with isinstance checks in _core.py and serde.py - Improve type safety by using explicit protocol type checks Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

- Revert `const_value` to `TensorProtocol | None` (no SparseTensorProtocol) - Add new `const_sparse_value: SparseTensorProtocol | None` field to ValueProtocol and Value - Update serde.py to use `const_sparse_value` for sparse initializers - Update tests accordingly Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

…se input support; suppress spurious warning for sparse-typed initializers Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

…enience function Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

…remove quoted type annotations Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

- Move scipy sparse conversion out of ir.tensor() into SparseTensor.from_scipy_sparse() so tensor() stays dense-only with a clean TensorProtocol return type. - Remove non-serializable doc_string/metadata_props from SparseTensor (SparseTensorProto only stores values/indices/dims); keep meta for in-memory analysis and document that it is not serialized. - Fix docs: const_sparse_value (not const_value) for sparse initializers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

SparseTensor.from_scipy_sparse() and numpy() require scipy. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Copilot

Pull request overview

Adds first-class sparse tensor support to the ONNX IR by introducing an in-memory SparseTensor representation (COO), enabling sparse (de)serialization and sparse initializers/attributes, and documenting + testing the new behavior.

Changes:

Introduces SparseTensor and SparseTensorProtocol, plus Value.const_sparse_value for sparse initializers/constants.
Adds sparse tensor (de)serialization in serde (including sparse initializer handling and sparse-tensor attributes).
Expands docs, public exports, tests, and CI deps to cover sparse tensor functionality.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/onnx_ir/serde.py	Adds sparse tensor (de)serialization, sparse initializer support, and sparse attribute handling.
src/onnx_ir/serde_test.py	Adds unit tests for sparse tensor serde, attributes, and sparse initializers.
src/onnx_ir/_protocols.py	Extends protocols with `SparseTensorProtocol` and `ValueProtocol.const_sparse_value`.
src/onnx_ir/_core.py	Adds `SparseTensor` implementation and `Value.const_sparse_value` plumbing.
src/onnx_ir/_core_test.py	Adds behavior tests for `SparseTensor.numpy()` / `as_tensor()` and scipy interop.
src/onnx_ir/_convenience/_constructors.py	Clarifies `tensor()` as dense-only and points users to `SparseTensor.from_scipy_sparse`.
src/onnx_ir/_convenience/_constructors_test.py	Adds tests for `SparseTensor.from_scipy_sparse()` construction/roundtrip.
src/onnx_ir/init.py	Exports `SparseTensor` from the public API.
noxfile.py	Adds `scipy` to common test dependencies.
docs/tensors.md	Documents sparse tensor usage, serde, initializers, and attributes.

coo_matrix only exposes .coords on scipy>=1.13; fall back to (row, col) for legacy spmatrix inputs. Addresses PR review feedback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

titaiwangms · 2026-06-11T21:33:22Z

Review summary — Support sparse tensors

Reviewed across readability, correctness, adversarial, spec-adherence, and cross-module integration lenses. One ship-blocking bug, confirmed independently with executable evidence (ONNX checker + numerical repros against onnx.in.proto).

🔴 Critical (blocking)

2-D sparse indices are transposed relative to the ONNX spec. The code uses [rank, NNZ]; SparseTensorProto requires [NNZ, rank] (each row is one nonzero's full coordinate tuple).

_core.py numpy() / as_tensor(): coords = tuple(indices_array[i] for i in range(indices_array.shape[0]))
from_scipy_sparse(): np.stack(coords, axis=0)

Impact: IR-created sparse tensors serialize to non-conformant protos; spec-compliant models with NNZ > rank crash (IndexError: too many indices), and asymmetric patterns silently densify to the wrong tensor. The test suite passes only because every 2-D case uses symmetric NNZ == rank == 2, where both layouts coincide.

Fix:

read: coords = tuple(indices_array[:, j] for j in range(indices_array.shape[1]))
write: np.stack(coords, axis=1)
update docstrings (the "[rank, NNZ]" wording) and tests; add an asymmetric NNZ != rank round-trip test validated against ONNX's own densification.

🟠 Major

const_sparse_value is not wired into existing consumers. The field was added to Value, but many if const_value is None sites were not updated, so sparse initializers are silently dropped or crash:
- register_initializer raises ValueError for sparse initializers (_core.py).
- _cloner.py copies const_value but not const_sparse_value → cloning drops sparse data.
- _graph_containers.py tensors() / tensor_items() / get_tensor() raise on sparse initializers.
- _safetensors/__init__.py assert tensor is not None fails for sparse initializers.
- journaling/_wrappers.py wraps const_value.fset but not const_sparse_value.fset → mutations bypass the journal.
- passes/common/_c_api_utils.py strips sparse initializers before the C API call.
- passes/common/initializer_deduplication.py logs false "invalid initializer" warnings.
Deserialized sparse initializers are typed SparseTensorType instead of TensorType, diverging from the dense initializer path and potentially altering model semantics (serde.py).
String sparse tensors densify with 0 instead of "". np.zeros(dtype=object) yields integer 0; the spec default for string tensors is the empty string (as_tensor).
No one-of invariant between const_value and const_sparse_value. Both can be set simultaneously; serialization silently prefers sparse and a stale dense constant can be dropped. Consider clearing the sibling in each setter, or validating against value.type.
Serialization mutates shared state. value.const_sparse_value.name = value.name (and SparseTensor.name → values.name aliasing) mutates shared objects during serialization; unsafe when a tensor is shared across values/attributes. Prefer patching the proto after serialization.

🟡 Minor

as_tensor(lazy=True) closure captures mutable self; snapshot dims/values/indices/dtype locally.
from_scipy_sparse does not canonicalize indices; the spec requires ascending order without duplicates (coo.sum_duplicates() + sort).
values shape docstring says [nnz, ...]; the spec defines values strictly as 1-D [NNZ].
Value.__repr__ labels a sparse backing as const_value={...}; should read const_sparse_value=.
Index-unpacking logic is duplicated verbatim in numpy() and as_tensor(); extract a single _unpack_indices helper.

✅ Praise

Name ↔ values.name mapping is spec-faithful across the property, Value.name rename propagation, and serialization.
1-D linear-index path (np.unravel_index, C order) is correct.
Sparse-vs-dense initializer routing and the targeted warning suppression are clean and well-tested.
SparseTensor deliberately avoids structurally satisfying the runtime-checkable TensorProtocol, so to_proto/from_proto dispatch stays correct.

_{🤖 Generated with multi-model review (readability, code, critical, deep-spec, integration).}

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 10 comments.

+        coords = getattr(coo, "coords", None)
+        if coords is None:
+            coords = (coo.row, coo.col)
+        indices_array = np.stack(coords, axis=0).astype(np.int64)
+        indices_tensor = Tensor(indices_array, _enums.DataType.INT64)
+        return cls(values_tensor, indices_tensor, list(array.shape))


+        values_array = self._values.numpy()
+        indices_array = self._indices.numpy()
+        if indices_array.ndim == 1:
+            # Linear (flat) indices → per-dimension coordinates
+            coords = np.unravel_index(indices_array, self._dims)
+        else:
+            # Shape [rank, NNZ] - each row is one dimension's indices
+            coords = tuple(indices_array[i] for i in range(indices_array.shape[0]))
+        return _scipy_sparse.coo_array((values_array, coords), shape=tuple(self._dims))


+            if indices_array.ndim == 1:
+                # Linear (flat) indices in row-major order
+                multi_indices = np.unravel_index(indices_array, self._dims)
+            else:
+                # Shape [rank, NNZ] - each row is one dimension's indices
+                multi_indices = tuple(indices_array[i] for i in range(indices_array.shape[0]))
+            dense[multi_indices] = values_array


+        values = _core.Tensor(np.array([1.0, 2.0], dtype=np.float32))
+        # Shape [rank=2, NNZ=2]: positions [0,1] and [1,0]
+        indices = _core.Tensor(np.array([[0, 1], [1, 0]], dtype=np.int64))
+        sparse = ir.SparseTensor(values=values, indices=indices, dims=[3, 3])
+        dense = sparse.as_tensor()
+        expected = np.zeros((3, 3), dtype=np.float32)
+        expected[0, 1] = 1.0
+        expected[1, 0] = 2.0
+        np.testing.assert_array_equal(dense.numpy(), expected)


+        values = _core.Tensor(np.array([1.0, 2.0], dtype=np.float32))
+        indices = _core.Tensor(np.array([[0, 1], [1, 0]], dtype=np.int64))
+        sparse = ir.SparseTensor(values=values, indices=indices, dims=[3, 3])
+        result = sparse.numpy()


+        data = np.array([1.0, 2.0], dtype=np.float32)
+        coo = sp.coo_array((data, ([0, 1], [1, 0])), shape=(3, 3))
+        sparse = ir.SparseTensor.from_scipy_sparse(coo, name="my_sparse")
+        self.assertIsInstance(sparse, ir.SparseTensor)
+        self.assertEqual(sparse.name, "my_sparse")
+        self.assertEqual(sparse.dims, [3, 3])
+        np.testing.assert_array_equal(sparse.values.numpy(), data)
+        np.testing.assert_array_equal(
+            sparse.indices.numpy(), np.array([[0, 1], [1, 0]], dtype=np.int64)
+        )


+        Returns:
+            A :class:`SparseTensor` in COO format with 2-D ``[rank, NNZ]`` indices.
+


+        * **2-D indices** (shape ``[rank, NNZ]``): per-dimension indices where
+          row *i* contains the index along dimension *i* for every non-zero.


+        * **2-D indices** (shape ``[rank, NNZ]``): per-dimension indices where
+          row *i* contains the index along dimension *i* for every non-zero.


+        self.assertEqual(list(result_proto.dims), [6])
+


justinchuby · 2026-06-27T00:09:59Z

@copilot please fix reviews

The ONNX SparseTensorProto spec requires 2-D indices in [NNZ, rank] layout (each row is one non-zero's full coordinate tuple), but the code was using [rank, NNZ] (each row is one dimension's indices). Bugs fixed: - from_scipy_sparse(): np.stack(coords, axis=0) -> axis=1 so shape is [NNZ, rank] - numpy(): row indexing indices_array[i] -> column indexing indices_array[:, j] - as_tensor(): same fix as numpy() - Docstrings updated throughout to say [NNZ, rank] instead of [rank, NNZ] The previous tests only used NNZ==rank==2, making both layouts identical and masking the bug. New asymmetric NNZ!=rank tests are added to _core_test.py, _constructors_test.py, and serde_test.py to definitively validate the correct index layout. Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

 @typing.overload
 def from_proto(proto: onnx.TensorProto) -> _protocols.TensorProtocol: ...  # type: ignore[overload-overlap]
 @typing.overload
+def from_proto(proto: onnx.SparseTensorProto) -> _core.SparseTensor: ...  # type: ignore[overload-overlap]


 @typing.overload
 def to_proto(ir_object: _protocols.TensorProtocol) -> onnx.TensorProto: ...  # type: ignore[overload-overlap]
 @typing.overload
+def to_proto(ir_object: _protocols.SparseTensorProtocol) -> onnx.SparseTensorProto: ...  # type: ignore[overload-overlap]


Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

+        if not isinstance(values, _protocols.TensorProtocol):
+            raise TypeError(
+                f"'values' must be a TensorProtocol, got {type(values)!r}. "
+                "To construct from a scipy sparse array, use SparseTensor.from_scipy_sparse()."
+            )
+        self._values: _protocols.TensorProtocol = values
+        self._indices: _protocols.TensorProtocol = indices
+        self._dims: list[int] = list(dims)


+        values_array = self._values.numpy()
+        indices_array = self._indices.numpy()
+        if indices_array.ndim == 1:
+            # Linear (flat) indices → per-dimension coordinates
+            coords = np.unravel_index(indices_array, self._dims)
+        else:
+            # Shape [NNZ, rank] - each row is one non-zero's coordinate tuple
+            coords = tuple(indices_array[:, j] for j in range(indices_array.shape[1]))
+        return _scipy_sparse.coo_array((values_array, coords), shape=tuple(self._dims))


+            values_array = self._values.numpy()
+            indices_array = self._indices.numpy()
+            dtype = self._values.dtype
+            np_dtype = dtype.numpy()
+            dense = np.zeros(self._dims, dtype=np_dtype)
+            if indices_array.ndim == 1:
+                # Linear (flat) indices in row-major order
+                multi_indices = np.unravel_index(indices_array, self._dims)
+            else:
+                # Shape [NNZ, rank] - each row is one non-zero's coordinate tuple
+                multi_indices = tuple(
+                    indices_array[:, j] for j in range(indices_array.shape[1])
+                )
+            dense[multi_indices] = values_array
+            return Tensor(dense, dtype)


justinchuby requested review from a team and titaiwangms as code owners June 11, 2026 19:00

github-code-quality Bot found potential problems Jun 11, 2026

View reviewed changes

Comment thread src/onnx_ir/serde.py Fixed

Comment thread src/onnx_ir/serde.py Fixed

justinchuby requested a review from Copilot June 11, 2026 20:59

justinchuby added module: api module: core labels Jun 11, 2026

Copilot started reviewing on behalf of justinchuby June 11, 2026 20:59 View session

justinchuby and others added 9 commits June 11, 2026 14:01

feat: Add SparseTensor.as_tensor(lazy=) method for dense conversion

e6fffd0

Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Move scipy sparse support from SparseTensor.__init__ to tensor() conv…

fbab907

…enience function Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Fix ruff lint errors: replace EN DASH with hyphen-minus in comments, …

993c7ab

…remove quoted type annotations Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Add scipy to CI test dependencies

4079045

SparseTensor.from_scipy_sparse() and numpy() require scipy. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

justinchuby force-pushed the justinchuby/sparse-tensor branch from f97f4d3 to 4079045 Compare June 11, 2026 21:01

Copilot AI reviewed Jun 11, 2026

View reviewed changes

Comment thread src/onnx_ir/_core.py

justinchuby added this to the 1.0.0 milestone Jun 26, 2026

justinchuby requested a review from Copilot June 26, 2026 21:29

Copilot started reviewing on behalf of justinchuby June 26, 2026 21:30 View session

Copilot AI reviewed Jun 26, 2026

View reviewed changes

Copilot started work on behalf of justinchuby June 27, 2026 00:10 View session

Copilot finished work on behalf of justinchuby June 27, 2026 00:17

github-code-quality Bot found potential problems Jun 27, 2026

View reviewed changes

justinchuby requested a review from Copilot June 27, 2026 01:50

Copilot started reviewing on behalf of justinchuby June 27, 2026 01:51 View session

Copilot AI reviewed Jun 27, 2026

View reviewed changes

justinchuby marked this pull request as draft August 8, 2026 00:33

Navigation Menu

Uh oh!

Support sparse tensors - #440

Support sparse tensors#440
justinchuby wants to merge 11 commits into
mainfrom
justinchuby/sparse-tensor

justinchuby commented Jun 11, 2026

Uh oh!

Uh oh!

Uh oh!

codecov Bot commented Jun 11, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

titaiwangms commented Jun 11, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

justinchuby commented Jun 27, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		Returns:
		A :class:`SparseTensor` in COO format with 2-D ``[rank, NNZ]`` indices.

		* 2-D indices (shape ``[rank, NNZ]``): per-dimension indices where
		row i contains the index along dimension i for every non-zero.

Uh oh!

Conversation

justinchuby commented Jun 11, 2026

Summary

Sparse tensor support

Design notes

Public API & docs

Testing

Uh oh!

Uh oh!

Uh oh!

codecov Bot commented Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

titaiwangms commented Jun 11, 2026

Review summary — Support sparse tensors

🔴 Critical (blocking)

🟠 Major

🟡 Minor

✅ Praise

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

justinchuby commented Jun 27, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

codecov Bot commented Jun 11, 2026 •

edited

Loading