BUG: np.take handle 64-bit indices on 32-bit platforms

See numpy#25607

Author: JuliaPoo

numpy/_core/src/multiarray/item_selection.c

@@ -231,21 +231,32 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
                  PyArrayObject *out, NPY_CLIPMODE clipmode)
 {
     PyArray_Descr *dtype;
-    PyArrayObject *obj = NULL, *self, *indices;
+    PyArrayObject *obj = NULL, *self, *indices, *tmp;
     npy_intp nd, i, n, m, max_item, chunk, itemsize, nelem;
     npy_intp shape[NPY_MAXDIMS];
 
     npy_bool needs_refcounting;
 
-    indices = NULL;
+    indices = tmp = NULL;
     self = (PyArrayObject *)PyArray_CheckAxis(self0, &axis,
                                     NPY_ARRAY_CARRAY_RO);
     if (self == NULL) {
         return NULL;
     }
-    indices = (PyArrayObject *)PyArray_ContiguousFromAny(indices0,
-                                                         NPY_INTP,
-                                                         0, 0);
+    tmp = (PyArrayObject *)PyArray_ContiguousFromAny(indices0,
+                                                     NPY_INT64,
+                                                     0, 0);
+    if (tmp == NULL) {
+        goto fail;
+    }
+
+    if (NPY_SIZEOF_INTP != 8) {
+        indices = (PyArrayObject *)PyArray_Cast(tmp, NPY_INTP);
+        Py_DECREF(tmp);
+    } else {
+        indices = tmp;
+    }
+
     if (indices == NULL) {
         goto fail;
     }

numpy/_core/tests/test_numeric.py

@@ -324,6 +324,38 @@ def test_take(self):
         out = np.take(a, indices)
         assert_equal(out, tgt)
 
+        # take 32 64
+        x32 = np.array([1, 2, 3, 4, 5], dtype=np.int32)
+        x64 = np.array([0, 2, 2, 3], dtype=np.int64)
+        tgt = np.array([1, 3, 3, 4], dtype=np.int32)
+        out = np.take(x32, x64)
+        assert_equal(out, tgt)
+        assert_equal(out.dtype, tgt.dtype)
+
+        # take 64 32
+        x64 = np.array([1, 2, 3, 4, 5], dtype=np.int64)
+        x32 = np.array([0, 2, 2, 3], dtype=np.int32)
+        tgt = np.array([1, 3, 3, 4], dtype=np.int64)
+        out = np.take(x64, x32)
+        assert_equal(out, tgt)
+        assert_equal(out.dtype, tgt.dtype)
+
+        # take 32 32
+        x32_0 = np.array([1, 2, 3, 4, 5], dtype=np.int32)
+        x32_1 = np.array([0, 2, 2, 3], dtype=np.int32)
+        tgt = np.array([1, 3, 3, 4], dtype=np.int32)
+        out = np.take(x32_0, x32_1)
+        assert_equal(out, tgt)
+        assert_equal(out.dtype, tgt.dtype)
+
+        # take 64 64
+        x64_0 = np.array([1, 2, 3, 4, 5], dtype=np.int64)
+        x64_1 = np.array([0, 2, 2, 3], dtype=np.int64)
+        tgt = np.array([1, 3, 3, 4], dtype=np.int64)
+        out = np.take(x64_0, x64_1)
+        assert_equal(out, tgt)
+        assert_equal(out.dtype, tgt.dtype)
+
     def test_trace(self):
         c = [[1, 2], [3, 4], [5, 6]]
         assert_equal(np.trace(c), 5)