Skip to content

fix: widen ravel_multi_index to int64 on overflow#16

Open
andrinr wants to merge 1 commit into
tum-pbs:mainfrom
andrinr:fix/ravel-multi-index-int32-overflow
Open

fix: widen ravel_multi_index to int64 on overflow#16
andrinr wants to merge 1 commit into
tum-pbs:mainfrom
andrinr:fix/ravel-multi-index-int32-overflow

Conversation

@andrinr

@andrinr andrinr commented May 20, 2026

Copy link
Copy Markdown

Summary

NumpyBackend.ravel_multi_index unconditionally cast the linearised result to the caller's multi_index.dtype. When that dtype was int32 and the shape product exceeded 2**31, the cast silently wrapped to negative. Downstream callers — primarily CompressedSparseMatrix.compress in phiml/math/_sparse.py — then surface it as:

ValueError: index <negative> is out of bounds for array with size <product>

from np.unravel_index on the next line.

Fix: preserve the caller's dtype when it fits, widen to int64 when np.prod(shape) > np.iinfo(dtype).max. The JAX backend already doesn't downcast, so the patch brings the numpy backend in line.

Reproducer

import numpy as np
from phiml.backend._numpy_backend import NUMPY

indices = np.array([[40000, 40000]], dtype=np.int32)
shape = (200_000, 200_000)   # product = 4e10 > 2**31

lin = NUMPY.ravel_multi_index(indices, shape)
# before this PR:
#   lin.dtype == int32, lin == [-589894592]
#   NUMPY.unravel_index(lin, shape)
#     → ValueError: index -589894592 is out of bounds for array with size 40000000000
# after this PR:
#   lin.dtype == int64, lin == [8000040000]
#   NUMPY.unravel_index(lin, shape) → [[40000, 40000]]   (round-trips)

In real use this surfaces in phi.physics.fluid.make_incompressible on a 3-D Navier-Stokes grid as soon as c_dims.volume * u_dims.volume > 2**31 — for an velocity-on-faces problem, that's around N ≥ 48.

What changed

  • phiml/backend/_numpy_backend.py:362-378 — drop the unconditional .astype(multi_index.dtype), pick int64 when overflow would occur.
  • tests/commit/backend/test__backend.py — new test_ravel_multi_index_int32_overflow regression case (large-shape round-trip + small-shape dtype preservation).

Testing done

  • New regression test passes.
  • tests/commit/backend/test__backend.py (43 cases) and tests/commit/math/test__sparse.py pass with the patch. One pre-existing failure on test_convert (JAX/DLPack GPU issue, unrelated to this PR) was confirmed to also fail on main.

Closes #15.

NumpyBackend.ravel_multi_index unconditionally cast the linearised
result to the caller's multi_index.dtype. When that dtype was int32
and the shape product exceeded 2**31, the cast silently wrapped to
negative and downstream callers (CompressedSparseMatrix.compress in
phiml/math/_sparse.py) surfaced it as

    ValueError: index <negative> is out of bounds for array with size <product>

from np.unravel_index.

Preserve the caller's dtype when it fits, but widen to int64 when the
linear index space exceeds np.iinfo(dtype).max. Mirrors the JAX
backend's behaviour, which never downcasts.

Closes tum-pbs#15.
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.

Sparse matrix index overflow: unravel_index fails with int32 wraparound when c_dims.volume * u_dims.volume > 2³¹

1 participant