numpy.flip (#155)

Author: Hardcode84

numba_dpcomp/numba_dpcomp/mlir/numpy/funcs.py

@@ -52,6 +52,12 @@ def linalg_index_impl(builder, dim):
     if isinstance(dim, int):
         return builder.linalg_index(dim)
 
+def _fix_axis(axis, num_dims):
+    if axis < 0:
+        axis = axis + num_dims
+    assert axis >= 0 and axis < num_dims
+    return axis
+
 def _array_reduce(builder, arg, axis, body, get_init_value):
     if axis is None:
         shape = arg.shape
@@ -68,9 +74,7 @@ def _array_reduce(builder, arg, axis, body, get_init_value):
     elif isinstance(axis, int):
         shape = arg.shape
         num_dims = len(shape)
-        if axis < 0:
-            axis += num_dims
-        assert axis >= 0 and axis < num_dims
+        axis = _fix_axis(axis, num_dims)
         iterators = [('reduction' if i == axis else 'parallel') for i in range(num_dims)]
         dims1 = ','.join(['d%s' % i for i in range(num_dims)])
         dims2 = ','.join(['d%s' % i for i in range(num_dims) if i != axis])
@@ -88,6 +92,24 @@ def sum_impl(builder, arg, axis=None):
     return _array_reduce(builder, arg, axis, lambda a, b: a + b, lambda b, t: 0)
 
 
+@register_func('numpy.flip', numpy.flip)
+def flip_impl(builder, arg, axis=None):
+    shape = arg.shape
+    num_dims = len(shape)
+    if axis is None:
+        axis = (True,) * num_dims
+    elif isinstance(axis, int):
+        axis = _fix_axis(axis, num_dims)
+        l = [False] * num_dims
+        l[axis] = True
+        axis = tuple(l)
+    else:
+        return
+
+    offsets = [0 if not axis[i] else shape[i] - 1 for i in range(num_dims)]
+    strides = [1 if not axis[i] else - 1 for i in range(num_dims)]
+    return builder.subview(arg, offsets, shape, strides)
+
 def _get_numpy_type(builder, dtype):
     types = [
         (builder.int8,  numpy.int8),

numba_dpcomp/numba_dpcomp/mlir/tests/test_numpy.py

@@ -300,6 +300,33 @@ def test_reduce_axis(py_func, arr):
     jit_func = njit(py_func)
     assert_equal(py_func(arr), jit_func(arr))
 
+@parametrize_function_variants("py_func", [
+    'lambda a: np.flip(a)',
+    ])
+@pytest.mark.parametrize("arr", [
+    np.array([1,2,3,4,5,6], dtype=np.int32),
+    np.array([[1,2,3],[4,5,6]], dtype=np.int32),
+    np.array([[[1,2,3],[4,5,6]]], dtype=np.int32),
+    ])
+def test_flip1(py_func, arr):
+    jit_func = njit(py_func)
+    assert_equal(py_func(arr), jit_func(arr))
+
+@parametrize_function_variants("py_func", [
+    'lambda a: np.flip(a, axis=0)',
+    'lambda a: np.flip(a, axis=1)',
+    'lambda a: np.flip(a, axis=-1)',
+    'lambda a: np.flip(a, axis=-2)',
+    ])
+@pytest.mark.parametrize("arr", [
+    np.array([[[1,2,3],[4,5,6]]], dtype=np.int32),
+    np.array([[[1,2,3],[4,5,6]]], dtype=np.float32),
+    ])
+@pytest.mark.xfail
+def test_flip2(py_func, arr):
+    jit_func = njit(py_func)
+    assert_equal(py_func(arr), jit_func(arr))
+
 def test_sum_add():
     def py_func(a, b):
         return np.add(a, b).sum()