Repository: numba/numba

Issue:
np.transpose() support
The binding for `np.transpose()` does not exist at present, it really should!


Discussion:

Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -25,7 +25,7 @@
                                     iternext_impl, impl_ret_borrowed,
                                     impl_ret_new_ref, impl_ret_untracked)
 from numba.typing import signature
-from numba.extending import register_jitable
+from numba.extending import register_jitable, overload
 from . import quicksort, mergesort, slicing
 
 
@@ -1481,6 +1481,19 @@ def array_transpose_vararg(context, builder, sig, args):
     return array_transpose_tuple(context, builder, new_sig, new_args)
 
 
+@overload(np.transpose)
+def numpy_transpose(a, axes=None):
+
+    if axes is None:
+        def np_transpose_impl(arr):
+            return arr.transpose()
+    else:
+        def np_transpose_impl(arr, axes=None):
+            return arr.transpose(axes)
+
+    return np_transpose_impl
+
+
 @lower_getattr(types.Array, 'T')
 def array_T(context, builder, typ, value):
     if typ.ndim <= 1:
@@ -1637,6 +1650,13 @@ def array_reshape_vararg(context, builder, sig, args):
     return array_reshape(context, builder, new_sig, new_args)
 
 
+@overload(np.reshape)
+def np_reshape(a, shape):
+    def np_reshape_impl(a, shape):
+        return a.reshape(shape)
+    return np_reshape_impl
+
+
 @lower_builtin('array.ravel', types.Array)
 def array_ravel(context, builder, sig, args):
     # Only support no argument version (default order='C')


Test patch:
diff --git a/numba/tests/test_array_manipulation.py b/numba/tests/test_array_manipulation.py
--- a/numba/tests/test_array_manipulation.py
+++ b/numba/tests/test_array_manipulation.py
@@ -1,12 +1,14 @@
 from __future__ import print_function
 
 from functools import partial
+from itertools import permutations
 import numba.unittest_support as unittest
 
 import numpy as np
 
 from numba.compiler import compile_isolated, Flags
 from numba import jit, types, from_dtype, errors, typeof
+from numba.errors import TypingError
 from .support import TestCase, MemoryLeakMixin, CompilationCache, tag
 
 enable_pyobj_flags = Flags()
@@ -16,10 +18,36 @@
 no_pyobj_flags.set('nrt')
 
 
+def from_generic(pyfuncs_to_use):
+    """Decorator for generic check functions.
+        Iterates over 'pyfuncs_to_use', calling 'func' with the iterated
+        item as first argument. Example:
+
+        @from_generic(numpy_array_reshape, array_reshape)
+        def check_only_shape(pyfunc, arr, shape, expected_shape):
+            # Only check Numba result to avoid Numpy bugs
+            self.memory_leak_setup()
+            got = generic_run(pyfunc, arr, shape)
+            self.assertEqual(got.shape, expected_shape)
+            self.assertEqual(got.size, arr.size)
+            del got
+            self.memory_leak_teardown()
+    """
+    def decorator(func):
+        def result(*args, **kwargs):
+            return (func(pyfunc, *args, **kwargs) for pyfunc in pyfuncs_to_use)
+        return result
+    return decorator
+
+
 def array_reshape(arr, newshape):
     return arr.reshape(newshape)
 
 
+def numpy_array_reshape(arr, newshape):
+    return np.reshape(arr, newshape)
+
+
 def flatten_array(a):
     return a.flatten()
 
@@ -40,6 +68,15 @@ def transpose_array(a):
     return a.transpose()
 
 
+def numpy_transpose_array(a):
+    return np.transpose(a)
+
+def numpy_transpose_array_axes_kwarg(arr, axes):
+    return np.transpose(arr, axes=axes)
+
+def array_transpose_axes(arr, axes):
+    return arr.transpose(axes)
+
 def squeeze_array(a):
     return a.squeeze()
 
@@ -65,6 +102,7 @@ def as_strided1(a):
     strides = (a.strides[0] // 2,) + a.strides[1:]
     return np.lib.stride_tricks.as_strided(a, strides=strides)
 
+
 def as_strided2(a):
     # Rolling window example as in https://github.com/numba/numba/issues/1884
     window = 3
@@ -101,42 +139,53 @@ def setUp(self):
 
     @tag('important')
     def test_array_reshape(self):
-        pyfunc = array_reshape
-        def run(arr, shape):
-            cres = self.ccache.compile(pyfunc, (typeof(arr), typeof(shape)))
+        pyfuncs_to_use = [array_reshape, numpy_array_reshape]
+
+        def generic_run(pyfunc, arr, shape):
+            cres = compile_isolated(pyfunc, (typeof(arr), typeof(shape)))
             return cres.entry_point(arr, shape)
-        def check(arr, shape):
+
+        @from_generic(pyfuncs_to_use)
+        def check(pyfunc, arr, shape):
             expected = pyfunc(arr, shape)
             self.memory_leak_setup()
-            got = run(arr, shape)
+            got = generic_run(pyfunc, arr, shape)
             self.assertPreciseEqual(got, expected)
             del got
             self.memory_leak_teardown()
-        def check_only_shape(arr, shape, expected_shape):
+
+        @from_generic(pyfuncs_to_use)
+        def check_only_shape(pyfunc, arr, shape, expected_shape):
             # Only check Numba result to avoid Numpy bugs
             self.memory_leak_setup()
-            got = run(arr, shape)
+            got = generic_run(pyfunc, arr, shape)
             self.assertEqual(got.shape, expected_shape)
             self.assertEqual(got.size, arr.size)
             del got
             self.memory_leak_teardown()
-        def check_err_shape(arr, shape):
+
+        @from_generic(pyfuncs_to_use)
+        def check_err_shape(pyfunc, arr, shape):
             with self.assertRaises(NotImplementedError) as raises:
-                run(arr, shape)
+                generic_run(pyfunc, arr, shape)
             self.assertEqual(str(raises.exception),
                              "incompatible shape for array")
-        def check_err_size(arr, shape):
+
+        @from_generic(pyfuncs_to_use)
+        def check_err_size(pyfunc, arr, shape):
             with self.assertRaises(ValueError) as raises:
-                run(arr, shape)
+                generic_run(pyfunc, arr, shape)
             self.assertEqual(str(raises.exception),
                              "total size of new array must be unchanged")
 
-        def check_err_multiple_negative(arr, shape):
+        @from_generic(pyfuncs_to_use)
+        def check_err_multiple_negative(pyfunc, arr, shape):
             with self.assertRaises(ValueError) as raises:
-                run(arr, shape)
+                generic_run(pyfunc, arr, shape)
             self.assertEqual(str(raises.exception),
                              "multiple negative shape values")
 
+
         # C-contiguous
         arr = np.arange(24)
         check(arr, (24,))
@@ -201,6 +250,82 @@ def check_empty(arr):
         # Exceptions leak references
         self.disable_leak_check()
 
+    def test_array_transpose_axes(self):
+        pyfuncs_to_use = [numpy_transpose_array_axes_kwarg,
+                          array_transpose_axes]
+
+        def run(pyfunc, arr, axes):
+            cres = self.ccache.compile(pyfunc, (typeof(arr), typeof(axes)))
+            return cres.entry_point(arr, axes)
+
+        @from_generic(pyfuncs_to_use)
+        def check(pyfunc, arr, axes):
+            expected = pyfunc(arr, axes)
+            got = run(pyfunc, arr, axes)
+            self.assertPreciseEqual(got, expected)
+            self.assertEqual(got.flags.f_contiguous,
+                             expected.flags.f_contiguous)
+            self.assertEqual(got.flags.c_contiguous,
+                             expected.flags.c_contiguous)
+
+        @from_generic(pyfuncs_to_use)
+        def check_err_axis_repeated(pyfunc, arr, axes):
+            with self.assertRaises(ValueError) as raises:
+                run(pyfunc, arr, axes)
+            self.assertEqual(str(raises.exception),
+                             "repeated axis in transpose")
+
+        @from_generic(pyfuncs_to_use)
+        def check_err_axis_oob(pyfunc, arr, axes):
+            with self.assertRaises(ValueError) as raises:
+                run(pyfunc, arr, axes)
+            self.assertEqual(str(raises.exception),
+                             "axis is out of bounds for array of given dimension")
+
+        @from_generic(pyfuncs_to_use)
+        def check_err_invalid_args(pyfunc, arr, axes):
+            with self.assertRaises((TypeError, TypingError)):
+                run(pyfunc, arr, axes)
+
+        arrs = [np.arange(24),
+                np.arange(24).reshape(4, 6),
+                np.arange(24).reshape(2, 3, 4),
+                np.arange(24).reshape(1, 2, 3, 4),
+                np.arange(64).reshape(8, 4, 2)[::3,::2,:]]
+
+        for i in range(len(arrs)):
+            # First check `None`, the default, which is to reverse dims
+            check(arrs[i], None)
+            # Check supplied axis permutations
+            for axes in permutations(tuple(range(arrs[i].ndim))):
+                ndim = len(axes)
+                neg_axes = tuple([x - ndim for x in axes])
+                check(arrs[i], axes)
+                check(arrs[i], neg_axes)
+
+        # Exceptions leak references
+        self.disable_leak_check()
+
+        check_err_invalid_args(arrs[1], "foo")
+        check_err_invalid_args(arrs[1], ("foo",))
+        check_err_invalid_args(arrs[1], 5.3)
+        check_err_invalid_args(arrs[2], (1.2, 5))
+
+        check_err_axis_repeated(arrs[1], (0, 0))
+        check_err_axis_repeated(arrs[2], (2, 0, 0))
+        check_err_axis_repeated(arrs[3], (3, 2, 1, 1))
+
+        check_err_axis_oob(arrs[0], (1,))
+        check_err_axis_oob(arrs[0], (-2,))
+        check_err_axis_oob(arrs[1], (0, 2))
+        check_err_axis_oob(arrs[1], (-3, 2))
+        check_err_axis_oob(arrs[1], (0, -3))
+        check_err_axis_oob(arrs[2], (3, 1, 2))
+        check_err_axis_oob(arrs[2], (-4, 1, 2))
+        check_err_axis_oob(arrs[3], (3, 1, 2, 5))
+        check_err_axis_oob(arrs[3], (3, 1, 2, -5))
+
+
     @tag('important')
     def test_expand_dims(self):
         pyfunc = expand_dims
@@ -391,16 +516,19 @@ def test_ravel_array_size_npm(self):
         self.test_ravel_array_size(flags=no_pyobj_flags)
 
     def test_transpose_array(self, flags=enable_pyobj_flags):
-        a = np.arange(9).reshape(3, 3)
+        @from_generic([transpose_array, numpy_transpose_array])
+        def check(pyfunc):
+            a = np.arange(9).reshape(3, 3)
 
-        pyfunc = transpose_array
-        arraytype1 = typeof(a)
-        cr = compile_isolated(pyfunc, (arraytype1,), flags=flags)
-        cfunc = cr.entry_point
+            arraytype1 = typeof(a)
+            cr = compile_isolated(pyfunc, (arraytype1,), flags=flags)
+            cfunc = cr.entry_point
 
-        expected = pyfunc(a)
-        got = cfunc(a)
-        np.testing.assert_equal(expected, got)
+            expected = pyfunc(a)
+            got = cfunc(a)
+            np.testing.assert_equal(expected, got)
+
+        check()
 
     def test_transpose_array_npm(self):
         self.test_transpose_array(flags=no_pyobj_flags)
diff --git a/numba/tests/test_array_methods.py b/numba/tests/test_array_methods.py
--- a/numba/tests/test_array_methods.py
+++ b/numba/tests/test_array_methods.py
@@ -64,9 +64,6 @@ def array_T(arr):
 def array_transpose(arr):
     return arr.transpose()
 
-def array_transpose_axes(arr, axes):
-    return arr.transpose(axes)
-
 def array_copy(arr):
     return arr.copy()
 
@@ -466,68 +463,6 @@ def check_arr(arr):
         arr = np.array([0]).reshape(())
         check_arr(arr)
 
-    def test_array_transpose_axes(self):
-        pyfunc = array_transpose_axes
-        def run(arr, axes):
-            cres = self.ccache.compile(pyfunc, (typeof(arr), typeof(axes)))
-            return cres.entry_point(arr, axes)
-        def check(arr, axes):
-            expected = pyfunc(arr, axes)
-            got = run(arr, axes)
-            self.assertPreciseEqual(got, expected)
-            self.assertEqual(got.flags.f_contiguous,
-                             expected.flags.f_contiguous)
-            self.assertEqual(got.flags.c_contiguous,
-                             expected.flags.c_contiguous)
-        def check_err_axis_repeated(arr, axes):
-            with self.assertRaises(ValueError) as raises:
-                run(arr, axes)
-            self.assertEqual(str(raises.exception),
-                             "repeated axis in transpose")
-        def check_err_axis_oob(arr, axes):
-            with self.assertRaises(ValueError) as raises:
-                run(arr, axes)
-            self.assertEqual(str(raises.exception),
-                             "axis is out of bounds for array of given dimension")
-        def check_err_invalid_args(arr, axes):
-            with self.assertRaises((TypeError, TypingError)):
-                run(arr, axes)
-
-        arrs = [np.arange(24),
-                np.arange(24).reshape(4, 6),
-                np.arange(24).reshape(2, 3, 4),
-                np.arange(24).reshape(1, 2, 3, 4),
-                np.arange(64).reshape(8, 4, 2)[::3,::2,:]]
-
-        for i in range(len(arrs)):
-            for axes in permutations(tuple(range(arrs[i].ndim))):
-                ndim = len(axes)
-                neg_axes = tuple([x - ndim for x in axes])
-                check(arrs[i], axes)
-                check(arrs[i], neg_axes)
-
-        # Exceptions leak references
-        self.disable_leak_check()
-
-        check_err_invalid_args(arrs[1], "foo")
-        check_err_invalid_args(arrs[1], ("foo",))
-        check_err_invalid_args(arrs[1], 5.3)
-        check_err_invalid_args(arrs[2], (1.2, 5))
-
-        check_err_axis_repeated(arrs[1], (0,0))
-        check_err_axis_repeated(arrs[2], (2,0,0))
-        check_err_axis_repeated(arrs[3], (3,2,1,1))
-
-        check_err_axis_oob(arrs[0], (1,))
-        check_err_axis_oob(arrs[0], (-2,))
-        check_err_axis_oob(arrs[1], (0,2))
-        check_err_axis_oob(arrs[1], (-3,2))
-        check_err_axis_oob(arrs[1], (0,-3))
-        check_err_axis_oob(arrs[2], (3,1,2))
-        check_err_axis_oob(arrs[2], (-4,1,2))
-        check_err_axis_oob(arrs[3], (3,1,2,5))
-        check_err_axis_oob(arrs[3], (3,1,2,-5))
-
 
     def test_array_transpose(self):
         self.check_layout_dependent_func(array_transpose)