From 116c1034392fc629bef29cf1c728a89ed055d9c1 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Fri, 11 Sep 2026 09:54:33 -0700 Subject: [PATCH] Fix vector/matrix product (@) silently overflowing for narrow value dtypes _major_matvec/_minor_matvec/_major_matmat/_minor_matmat allocated their output accumulator as values.dtype, the compact per-element storage dtype (e.g. uint16 for single-cell counts), instead of a dtype sized for the accumulated total. Contributions landing in the same output slot wrapped modulo 2**bits instead of promoting, unlike sum() which already accumulates in float64. Promote values and the other operand to their common numpy dtype (np.result_type) before running the kernels, matching the dtype promotion an equivalent dense-array product would get. Fixes #43. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_ops.py | 19 +++++++++++++++---- tests/test_ops.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/src/vsparse/_ops.py b/src/vsparse/_ops.py index e9d3792..594fcbc 100644 --- a/src/vsparse/_ops.py +++ b/src/vsparse/_ops.py @@ -167,21 +167,32 @@ def minor_counts(value_ptr, indices, n_minor, nthreads): return partial.sum(axis=0) +def _promote(values, other): + """Widen ``values``/``other`` to their common dtype so accumulation doesn't + silently wrap modulo a narrow stored dtype (e.g. ``uint16`` counts), matching + the dtype promotion ordinary ``numpy`` arrays would get for the same product. + """ + out_dtype = np.result_type(values.dtype, other.dtype) + return values.astype(out_dtype, copy=False), other.astype(out_dtype, copy=False) + + def major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor): - return _major_matvec(major_ptr, values, value_ptr, indices, np.asarray(x), n_major, n_minor) + values, x = _promote(values, np.asarray(x)) + return _major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor) def minor_matvec(major_ptr, values, value_ptr, indices, x, n_major): - return _minor_matvec(major_ptr, values, value_ptr, indices, np.asarray(x), n_major) + values, x = _promote(values, np.asarray(x)) + return _minor_matvec(major_ptr, values, value_ptr, indices, x, n_major) def major_matmat(major_ptr, values, value_ptr, indices, b, n_major, n_minor): - b = np.ascontiguousarray(b) + values, b = _promote(values, np.ascontiguousarray(b)) return _major_matmat(major_ptr, values, value_ptr, indices, b, n_major, n_minor) def minor_matmat(major_ptr, values, value_ptr, indices, b, n_major): - b = np.ascontiguousarray(b) + values, b = _promote(values, np.ascontiguousarray(b)) return _minor_matmat(major_ptr, values, value_ptr, indices, b, n_major) diff --git a/tests/test_ops.py b/tests/test_ops.py index 04df920..487a545 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -135,3 +135,32 @@ def test_unsupported_matmul_operands_raise(dense, vcls): with pytest.raises(TypeError): _ = arr_3d @ v + +def test_matmul_does_not_overflow_narrow_value_dtype(vcls): + """Regression test for #43: matvec/matmat must not wrap modulo a narrow + stored dtype (e.g. uint16 counts) the way ``values.dtype``-sized + accumulators used to. ``sum()`` already got this right; the products + should agree with it instead of silently wrapping. + """ + n_rows, n_cols = 2000, 5 + dense = np.full((n_rows, n_cols), 40, dtype=np.uint16) # col totals = 80000 > uint16 max + v = _make(vcls, dense) + assert v.dtype == np.uint16 + + ones_rows = np.ones(n_rows, dtype=np.float64) + expected_cols = np.ravel(v.sum(axis=0)) + np.testing.assert_array_equal(ones_rows @ v, expected_cols) + + ones_mat = np.ones((3, n_rows), dtype=np.float64) + np.testing.assert_array_equal(ones_mat @ v, np.tile(expected_cols, (3, 1))) + + dense2 = np.zeros((3, 2000), dtype=np.uint16) + dense2[0, :] = 40 # row 0 total = 80000 > uint16 max + v2 = _make(vcls, dense2) + ones_cols = np.ones(2000, dtype=np.float64) + expected_rows = np.ravel(v2.sum(axis=1)) + np.testing.assert_array_equal(v2 @ ones_cols, expected_rows) + + ones_mat2 = np.ones((2000, 3), dtype=np.float64) + np.testing.assert_array_equal(v2 @ ones_mat2, np.tile(expected_rows, (3, 1)).T) +