Merge pull request numpy#26610 from JuliaPoo/issue-25607-take-32bit-bug

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

Author: ngoldbaum

numpy/_core/src/multiarray/arrayobject.h

@@ -55,6 +55,14 @@ static const int NPY_ARRAY_WAS_PYTHON_COMPLEX = (1 << 28);
 static const int NPY_ARRAY_WAS_INT_AND_REPLACED = (1 << 27);
 static const int NPY_ARRAY_WAS_PYTHON_LITERAL = (1 << 30 | 1 << 29 | 1 << 28);
 
+/*
+ * This flag allows same kind casting, similar to NPY_ARRAY_FORCECAST.
+ *
+ * An array never has this flag set; they're only used as parameter
+ * flags to the various FromAny functions.
+ */
+static const int NPY_ARRAY_SAME_KIND_CASTING = (1 << 26);
+
 #ifdef __cplusplus
 }
 #endif

numpy/_core/src/multiarray/ctors.c

@@ -7,6 +7,7 @@
 #include <structmember.h>
 
 #include "numpy/arrayobject.h"
+#include "arrayobject.h"
 #include "numpy/arrayscalars.h"
 
 #include "numpy/npy_math.h"
@@ -1908,6 +1909,10 @@ PyArray_FromArray(PyArrayObject *arr, PyArray_Descr *newtype, int flags)
         newtype->elsize = oldtype->elsize;
     }
 
+    if (flags & NPY_ARRAY_SAME_KIND_CASTING) {
+        casting = NPY_SAME_KIND_CASTING;
+    }
+
     /* If the casting if forced, use the 'unsafe' casting rule */
     if (flags & NPY_ARRAY_FORCECAST) {
         casting = NPY_UNSAFE_CASTING;

numpy/_core/src/multiarray/item_selection.c

@@ -243,9 +243,12 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
     if (self == NULL) {
         return NULL;
     }
-    indices = (PyArrayObject *)PyArray_ContiguousFromAny(indices0,
-                                                         NPY_INTP,
-                                                         0, 0);
+
+    indices = (PyArrayObject *)PyArray_FromAny(indices0,
+                PyArray_DescrFromType(NPY_INTP),
+                0, 0,
+                NPY_ARRAY_SAME_KIND_CASTING | NPY_ARRAY_DEFAULT,
+                NULL);
     if (indices == NULL) {
         goto fail;
     }

numpy/_core/tests/test_numeric.py

@@ -324,6 +324,18 @@ def test_take(self):
         out = np.take(a, indices)
         assert_equal(out, tgt)
 
+        pairs = [
+            (np.int32, np.int32), (np.int32, np.int64),
+            (np.int64, np.int32), (np.int64, np.int64)
+        ]
+        for array_type, indices_type in pairs:
+            x = np.array([1, 2, 3, 4, 5], dtype=array_type)
+            ind = np.array([0, 2, 2, 3], dtype=indices_type)
+            tgt = np.array([1, 3, 3, 4], dtype=array_type)
+            out = np.take(x, ind)
+            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)