Repository: numba/numba

Issue:
Support axes argument in transpose
One often needs to convert images (stored in ndarrays) between formats like HWC (height/width/channel) and CHW, when working with deep learning frameworks like TensorFlow and PyTorch. This requires permutations of axes which transpose provides with the optional `axes` argument.


Discussion:
@agayev

The output array of `transpose` is possibly neither `C` layout nor `F` layout if we return a view. Can Numba handle this? @sklam @stuartarchibald 
We might need to copy the data. @agayev
@ehsantn, Yes, numba has a 'A' for "any" layout for arbitrarily strided array.
@sklam Ok, thanks. I thought 'A' layout has to be resolved to 'C' or 'F' after type inference.

We assume 'C' layout in ParallelAccelerator which needs to be fixed.

Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -1382,10 +1382,105 @@ def src_cleanup():
 #-------------------------------------------------------------------------------
 # Shape / layout altering
 
+def vararg_to_tuple(context, builder, sig, args):
+    aryty = sig.args[0]
+    dimtys = sig.args[1:]
+    # values
+    ary = args[0]
+    dims = args[1:]
+    # coerce all types to intp
+    dims = [context.cast(builder, val, ty, types.intp)
+            for ty, val in zip(dimtys, dims)]
+    # make a tuple
+    shape = cgutils.pack_array(builder, dims, dims[0].type)
+
+    shapety = types.UniTuple(dtype=types.intp, count=len(dims))
+    new_sig = typing.signature(sig.return_type, aryty, shapety)
+    new_args = ary, shape
+
+    return new_sig, new_args
+
+
 @lower_builtin('array.transpose', types.Array)
 def array_transpose(context, builder, sig, args):
     return array_T(context, builder, sig.args[0], args[0])
 
+
+def permute_arrays(axis, shape, strides):
+    if len(axis) != len(set(axis)):
+        raise ValueError("repeated axis in transpose")
+    dim = len(shape)
+    for x in axis:
+        if x >= dim or abs(x) > dim:
+            raise ValueError("axis is out of bounds for array of given dimension")
+
+    shape[:] = shape[axis]
+    strides[:] = strides[axis]
+
+
+# Transposing an array involves permuting the shape and strides of the array
+# based on the given axes.
+@lower_builtin('array.transpose', types.Array, types.BaseTuple)
+def array_transpose_tuple(context, builder, sig, args):
+    aryty = sig.args[0]
+    ary = make_array(aryty)(context, builder, args[0])
+
+    axisty, axis = sig.args[1], args[1]
+    num_axis, dtype = axisty.count, axisty.dtype
+
+    ll_intp = context.get_value_type(types.intp)
+    ll_ary_size = lc.Type.array(ll_intp, num_axis)
+
+    # Allocate memory for axes, shapes, and strides arrays.
+    arys = [axis, ary.shape, ary.strides]
+    ll_arys = [cgutils.alloca_once(builder, ll_ary_size) for _ in arys]
+
+    # Store axes, shapes, and strides arrays to the allocated memory.
+    for src, dst in zip(arys, ll_arys):
+        builder.store(src, dst)
+
+    np_ary_ty = types.Array(dtype=dtype, ndim=1, layout='C')
+    np_itemsize = context.get_constant(types.intp,
+                                       context.get_abi_sizeof(ll_intp))
+
+    # Form NumPy arrays for axes, shapes, and strides arrays.
+    np_arys = [make_array(np_ary_ty)(context, builder) for _ in arys]
+
+    # Roughly, `np_ary = np.array(ll_ary)` for each of axes, shapes, and strides.
+    for np_ary, ll_ary in zip(np_arys, ll_arys):
+        populate_array(np_ary,
+                       data=builder.bitcast(ll_ary, ll_intp.as_pointer()),
+                       shape=[context.get_constant(types.intp, num_axis)],
+                       strides=[np_itemsize],
+                       itemsize=np_itemsize,
+                       meminfo=None)
+
+    # Pass NumPy arrays formed above to permute_arrays function that permutes
+    # shapes and strides based on axis contents.
+    context.compile_internal(builder, permute_arrays,
+                             typing.signature(types.void,
+                                              np_ary_ty, np_ary_ty, np_ary_ty),
+                             [a._getvalue() for a in np_arys])
+
+    # Make a new array based on permuted shape and strides and return it.
+    ret = make_array(sig.return_type)(context, builder)
+    populate_array(ret,
+                   data=ary.data,
+                   shape=builder.load(ll_arys[1]),
+                   strides=builder.load(ll_arys[2]),
+                   itemsize=ary.itemsize,
+                   meminfo=ary.meminfo,
+                   parent=ary.parent)
+    res = ret._getvalue()
+    return impl_ret_borrowed(context, builder, sig.return_type, res)
+
+
+@lower_builtin('array.transpose', types.Array, types.VarArg(types.Any))
+def array_transpose_vararg(context, builder, sig, args):
+    new_sig, new_args = vararg_to_tuple(context, builder, sig, args)
+    return array_transpose_tuple(context, builder, new_sig, new_args)
+
+
 @lower_getattr(types.Array, 'T')
 def array_T(context, builder, typ, value):
     if typ.ndim <= 1:
@@ -1538,21 +1633,7 @@ def array_reshape(context, builder, sig, args):
 
 @lower_builtin('array.reshape', types.Array, types.VarArg(types.Any))
 def array_reshape_vararg(context, builder, sig, args):
-    # types
-    aryty = sig.args[0]
-    dimtys = sig.args[1:]
-    # values
-    ary = args[0]
-    dims = args[1:]
-    # coerce all types to intp
-    dims = [context.cast(builder, val, ty, types.intp)
-            for ty, val in zip(dimtys, dims)]
-    # make a tuple
-    shape = cgutils.pack_array(builder, dims, dims[0].type)
-
-    shapety = types.UniTuple(dtype=types.intp, count=len(dims))
-    new_sig = typing.signature(sig.return_type, aryty, shapety)
-    new_args = ary, shape
+    new_sig, new_args = vararg_to_tuple(context, builder, sig, args)
     return array_reshape(context, builder, new_sig, new_args)
 
 
diff --git a/numba/typing/arraydecl.py b/numba/typing/arraydecl.py
--- a/numba/typing/arraydecl.py
+++ b/numba/typing/arraydecl.py
@@ -292,9 +292,39 @@ def _resolve_real_imag(self, ary, attr):
 
     @bound_function("array.transpose")
     def resolve_transpose(self, ary, args, kws):
-        assert not args
+        def sentry_shape_scalar(ty):
+            if ty in types.number_domain:
+                # Guard against non integer type
+                if not isinstance(ty, types.Integer):
+                    raise TypeError("transpose() arg cannot be {0}".format(ty))
+                return True
+            else:
+                return False
+
         assert not kws
-        return signature(self.resolve_T(ary))
+        if len(args) == 0:
+            return signature(self.resolve_T(ary))
+
+        if len(args) == 1:
+            shape, = args
+
+            if sentry_shape_scalar(shape):
+                assert ary.ndim == 1
+                return signature(ary, *args)
+
+            shape = normalize_shape(shape)
+            if shape is None:
+                return
+
+            assert ary.ndim == shape.count
+            return signature(self.resolve_T(ary), shape)
+
+        else:
+            if any(not sentry_shape_scalar(a) for a in args):
+                raise TypeError("transpose({0}) is not supported".format(
+                    ', '.join(args)))
+            assert ary.ndim == len(args)
+            return signature(self.resolve_T(ary), *args)
 
     @bound_function("array.copy")
     def resolve_copy(self, ary, args, kws):


Test patch:
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
@@ -1,6 +1,6 @@
 from __future__ import division
 
-from itertools import product, cycle
+from itertools import product, cycle, permutations
 import sys
 
 import numpy as np
@@ -64,6 +64,9 @@ 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()
 
@@ -456,6 +459,69 @@ 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)