numba / numba

NumPy aware dynamic Python compiler using LLVM
https://numba.pydata.org/
BSD 2-Clause "Simplified" License
9.95k stars 1.13k forks source link

Feature request: `numpy.flip` #4321

Closed jriehl closed 4 years ago

jriehl commented 5 years ago

Feature request

Implement numpy.flip() in arrayobj.py. This is a bit more difficult than just adding the Numpy Python code since it needs to work across n-dimensional arrays. I would considering using np.nditer or similar to acquire this functionality. Here is 1-d case from the sprint at SciPy2019:

def flip_ud(input_array):
  '''
  Quick implementaiton of flip_ud from
  Numpy to work with Numba
  '''
  new_array = np.empty(input_array.shape, dtype=input_array.dtype)
  arr_len = input_array.shape[0]
  for i in range(arr_len):
    new_array[arr_len - i - 1] = input_array[i]

  return new_array

The use case for this code came from a user trying to use np.flip() in a heuristic approach to efficiently solving the traveling salesman problem.

stuartarchibald commented 5 years ago

Thanks for the request. Here's some code that does it, this will hopefully start off a PR by someone :) It needs neatening and testing more:

from numba import njit
from numba.typed import List
from numba import types, typing
from numba.unsafe.tuple import tuple_setitem
from numba.extending import overload, intrinsic
import numpy as np

@overload(np.fliplr)
def np_flip_lr(a):
    def impl(a):
        return a[::, ::-1, ...]
    return impl

@overload(np.flipud)
def np_flip_ud(a):
    def impl(a):
        return a[::-1, ...]
    return impl

@intrinsic
def build_slice_tuple(tyctx, sz):
    """ Creates a tuple of slices for np.flip indexing like
    `(slice(None, None, -1),) * sz` """
    size = int(sz.literal_value)
    tuple_type = types.UniTuple(dtype=types.slice3_type, count=size)
    sig = tuple_type(sz)

    def codegen(context, builder, signature, args):
        def impl(length, empty_tuple):
            out = empty_tuple
            for i in range(length):
                out = tuple_setitem(out, i, slice(None, None, -1))
            return out

        inner_argtypes = [types.intp, tuple_type]
        inner_sig = typing.signature(tuple_type, *inner_argtypes)
        ll_idx_type = context.get_value_type(types.intp)
        # Allocate an empty tuple
        empty_tuple = context.get_constant_undef(tuple_type)
        inner_args = [ll_idx_type(size), empty_tuple]

        res = context.compile_internal(builder, impl, inner_sig, inner_args)
        return res

    return sig, codegen

@overload(np.flip)
def np_flip(a):
    ndim = a.ndim

    def impl(a):
        sl = build_slice_tuple(ndim)
        return a[sl]
    return impl

@njit
def test_flip(x):
    return np.flip(x), np.flipud(x), np.fliplr(x)

z = np.arange(20).reshape((4, 5))

np.testing.assert_allclose(test_flip.py_func(z), test_flip(z))

print("NumPy")
[print(x, "\n") for x in test_flip.py_func(z)]
print("Numba")
[print(x, "\n") for x in test_flip(z)]

# something more involved
shape = tuple([*range(1, 7)])
z = np.arange(np.prod(shape)).reshape(shape)
np.testing.assert_allclose(test_flip.py_func(z), test_flip(z))
rteja1113 commented 5 years ago

Hi @stuartarchibald , I'd like to work on this if you don't mind

stuartarchibald commented 5 years ago

@rteja1113 great, thanks for taking this on.

The use of Numba's extension API @overload decorator is strongly recommended for this task. A guide to using @overload is here and API documentation is here.

It may be useful to take a look at examples of pull requests adding support for a NumPy function that have already been done if guidance is needed... here's a couple:

LucaCappelletti94 commented 4 years ago

Has this been implemented yet?

stuartarchibald commented 4 years ago

Has this been implemented yet?

Yes, https://github.com/numba/numba/pull/5313, thanks for prompting, closing this as it's resolved.