BUG: np.take cannot deal with 64-bit indices on 32-bit platforms

[rgommers]

This looks like a bug, because regular indexing is able to deal with the exact same situation. Running these test cases, the ones where the second argument has int64 dtype fail in a 32-bit Linux Docker container:

def test_take_32_64():
    x32 = np.ones(2, dtype=np.int32)
    x64 = np.ones(2, dtype=np.int64)
    np.take(x32, x64)

def test_take_64_32():
    x32 = np.ones(2, dtype=np.int32)
    x64 = np.ones(2, dtype=np.int64)
    np.take(x64, x32)

def test_take_32_32():
    x32 = np.ones(2, dtype=np.int32)
    np.take(x32, x32)

def test_take_64_64():
    x64 = np.ones(2, dtype=np.int64)
    np.take(x64, x64)

Full tracebacks:

_______________________________ test_take_32_64 ________________________________
/scipy/test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:57: in _wrapfunc
    return bound(*args, **kwds)
E   TypeError: Cannot cast array data from dtype('int64') to dtype('int32') according to the rule 'safe'
        args       = (array([1, 1], dtype=int64),)
        bound      = <built-in method take of numpy.ndarray object at 0xe28f8aa0>
        kwds       = {'axis': None, 'mode': 'raise', 'out': None}
        method     = 'take'
        obj        = array([1, 1])
During handling of the above exception, another exception occurred:
scipy/_lib/tests/test__util.py:25: in test_take_32_64
    np.take(x32, x64)
        x32        = array([1, 1])
        x64        = array([1, 1], dtype=int64)
<__array_function__ internals>:180: in take
    ???
        args       = (array([1, 1]), array([1, 1], dtype=int64))
        dispatcher = <function _take_dispatcher at 0xeed66a00>
        implementation = <function take at 0xeed66a90>
        kwargs     = {}
        public_api = <function take at 0xeed66ad8>
        relevant_args = (array([1, 1]), None)
/scipy/test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:190: in take
    return _wrapfunc(a, 'take', indices, axis=axis, out=out, mode=mode)
        a          = array([1, 1])
        axis       = None
        indices    = array([1, 1], dtype=int64)
        mode       = 'raise'
        out        = None
/scipy/test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:66: in _wrapfunc
    return _wrapit(obj, method, *args, **kwds)
        args       = (array([1, 1], dtype=int64),)
        bound      = <built-in method take of numpy.ndarray object at 0xe28f8aa0>
        kwds       = {'axis': None, 'mode': 'raise', 'out': None}
        method     = 'take'
        obj        = array([1, 1])
/scipy/test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:43: in _wrapit
    result = getattr(asarray(obj), method)(*args, **kwds)
E   TypeError: Cannot cast array data from dtype('int64') to dtype('int32') according to the rule 'safe'
        args       = (array([1, 1], dtype=int64),)
        kwds       = {'axis': None, 'mode': 'raise', 'out': None}
        method     = 'take'
        obj        = array([1, 1])
        wrap       = <built-in method __array_wrap__ of numpy.ndarray object at 0xe28f8aa0>

_______________________________ test_take_64_64 ________________________________
/scipy/test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:57: in _wrapfunc
    return bound(*args, **kwds)
E   TypeError: Cannot cast array data from dtype('int64') to dtype('int32') according to the rule 'safe'
        args       = (array([1, 1], dtype=int64),)
        bound      = <built-in method take of numpy.ndarray object at 0xdff33f20>
        kwds       = {'axis': None, 'mode': 'raise', 'out': None}
        method     = 'take'
        obj        = array([1, 1], dtype=int64)
During handling of the above exception, another exception occurred:
scipy/_lib/tests/test__util.py:38: in test_take_64_64
    np.take(x64, x64)
        x64        = array([1, 1], dtype=int64)
<__array_function__ internals>:180: in take
    ???
        args       = (array([1, 1], dtype=int64), array([1, 1], dtype=int64))
        dispatcher = <function _take_dispatcher at 0xeed66a00>
        implementation = <function take at 0xeed66a90>
        kwargs     = {}
        public_api = <function take at 0xeed66ad8>
        relevant_args = (array([1, 1], dtype=int64), None)
/scipy/test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:190: in take
    return _wrapfunc(a, 'take', indices, axis=axis, out=out, mode=mode)
        a          = array([1, 1], dtype=int64)
        axis       = None
        indices    = array([1, 1], dtype=int64)
        mode       = 'raise'
        out        = None
/scipy/test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:66: in _wrapfunc
    return _wrapit(obj, method, *args, **kwds)
        args       = (array([1, 1], dtype=int64),)
        bound      = <built-in method take of numpy.ndarray object at 0xdff33f20>
        kwds       = {'axis': None, 'mode': 'raise', 'out': None}
        method     = 'take'
        obj        = array([1, 1], dtype=int64)
/scipy/test/lib/python3.9/site-packages/numpy/core/fromnumeric.py:43: in _wrapit
    result = getattr(asarray(obj), method)(*args, **kwds)
E   TypeError: Cannot cast array data from dtype('int64') to dtype('int32') according to the rule 'safe'
        args       = (array([1, 1], dtype=int64),)
        kwds       = {'axis': None, 'mode': 'raise', 'out': None}
        method     = 'take'
        obj        = array([1, 1], dtype=int64)
        wrap       = <built-in method __array_wrap__ of numpy.ndarray object at 0xdff33f20>
=========================== short test summary info ============================
FAILED scipy/_lib/tests/test__util.py::test_take_32_64 - TypeError: Cannot ca...
FAILED scipy/_lib/tests/test__util.py::test_take_64_64 - TypeError: Cannot ca..

It's expected that np.take has parity here with indexing, and accepts int64 dtype input also on 32-bit platforms. I came across this in SciPy code, where trying to replace indexing with take usage failed unexpectedly. It was code like this:

def _kpoints(data, k, rng):
    idx = rng.choice(data.shape[0], size=int(k), replace=False)
    return np.take(data, idx, axis=0)

There, idx is always 64-bit, and this is a fairly common way of generating indices, so it's cumbersome to have to cast idx to intp dtype.

[seberg]

Indexing just uses same-kind casting and accepts that you might get wrong results for strange values. That would be an OK change here, I think.