Closed jriehl closed 4 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))
Hi @stuartarchibald , I'd like to work on this if you don't mind
@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:
Has this been implemented yet?
Has this been implemented yet?
Yes, https://github.com/numba/numba/pull/5313, thanks for prompting, closing this as it's resolved.
Feature request
Implement
numpy.flip()
inarrayobj.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 usingnp.nditer
or similar to acquire this functionality. Here is 1-d case from the sprint at SciPy2019: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.