BUG: sparse: fix selecting wrong dtype for coo coords (scipy#22353)

* Fixes bug with selecting wrong dtype for coo coords

* add test for reshape having wrong dtype

* test that smaller dtype is maintained even if intermediate values are big

---------

Co-authored-by: Dan Schult <dschult@colgate.edu>

Author: 2 people

scipy/sparse/_coo.py

@@ -148,7 +148,7 @@ def reshape(self, *args, **kwargs):
         else:
             new_coords = np.unravel_index(flat_coords, shape, order=order)
 
-        idx_dtype = self._get_index_dtype(self.coords, maxval=max(self.shape))
+        idx_dtype = self._get_index_dtype(self.coords, maxval=max(shape))
         new_coords = tuple(np.asarray(co, dtype=idx_dtype) for co in new_coords)
 
         # Handle copy here rather than passing on to the constructor so that no

scipy/sparse/tests/test_coo.py

@@ -129,6 +129,22 @@ def test_non_subscriptability():
                        match="'coo_array' object is not subscriptable"):
         coo_2d[0, :]
 
+def test_reshape_overflow():
+    # see gh-22353 : new idx_dtype can need to be int64 instead of int32
+    M, N = (1045507, 523266)
+    coords = (np.array([M - 1], dtype='int32'), np.array([N - 1], dtype='int32'))
+    A = coo_array(([3.3], coords), shape=(M, N))
+
+    # need new idx_dtype to not overflow
+    B = A.reshape((M * N, 1))
+    assert B.coords[0].dtype == np.dtype('int64')
+    assert B.coords[0][0] == (M * N) - 1
+
+    # need idx_dtype to stay int32 if before and after can be int32
+    C = A.reshape(N, M)
+    assert C.coords[0].dtype == np.dtype('int32')
+    assert C.coords[0][0] == N - 1
+
 def test_reshape():
     arr1d = coo_array([1, 0, 3])
     assert arr1d.shape == (3,)