Use np.transpose with axes argument for arbitrary array dimension

Numba Support: How do I do ...?
1

Hi there,

I am trying to use np.transpose with the axes argument. It seems that this argument needs to be a tuple and not a list. Hence, I am trying to create dynamically a tuple.

import numba
import numpy as np

# just an example
arr = np.arange(8).reshape(2, 2, 2)
axes = (1, 0, *range(2, arr.ndim))

# this works fine
@numba.njit
def transpose1(a, axes):
    return np.transpose(a, axes)

@numba.njit
def transpose2(a):
    axes = (1, 0, *range(2, a.ndim))
    return np.transpose(a, axes)

@numba.njit
def transpose3(a):
    return np.transpose(a, (1, 0))

Using above definitions results in following output

>>> transpose1(arr, axes)
array([[[0, 1],
        [4, 5]],

       [[2, 3],
        [6, 7]]])

# but this not anymore
>>> transpose2(arr)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/braniii/Documents/Coding/numba/numba/core/dispatcher.py", line 421, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/braniii/Documents/Coding/numba/numba/core/dispatcher.py", line 362, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function add>) found for signature:
 
 >>> add(Tuple(Literal[int](1), Literal[int](0)), range_state_int64)
 
There are 18 candidate implementations:
  - Of which 16 did not match due to:
  Overload of function 'add': File: <numerous>: Line N/A.
    With argument(s): '(UniTuple(int64 x 2), range_state_int64)':
   No match.
  - Of which 2 did not match due to:
  Operator Overload in function 'add': File: unknown: Line unknown.
    With argument(s): '(UniTuple(int64 x 2), range_state_int64)':
   No match for registered cases:
    * (int64, int64) -> int64
    * (int64, uint64) -> int64
    * (uint64, int64) -> int64
    * (uint64, uint64) -> uint64
    * (float32, float32) -> float32
    * (float64, float64) -> float64
    * (complex64, complex64) -> complex64
    * (complex128, complex128) -> complex128

During: typing of intrinsic-call at <stdin> (3)

File "<stdin>", line 3:
<source missing, REPL/exec in use?>

# using a tuple with only two keys does not work either
# works for
>>> transpose3(arr[:, :, 0])
array([[0, 4],
       [2, 6]])

# but fails for ndim != 2
>>> transpose3(arr)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 2, in transpose
  File "<__array_function__ internals>", line 5, in transpose
  File "/home/braniii/.local/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 658, in transpose
    return _wrapfunc(a, 'transpose', axes)
  File "/home/braniii/.local/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 58, in _wrapfunc
    return bound(*args, **kwds)
ValueError: axes don't match array

My problem is, that I need to call np.transpose to implement a needed (from my side) numpy function. So using python mode is not an option. Has someone an idea? Thx.

Daniel

2

Hi @braniii ,

it might be worth adding the error messages that you get, since this will make it easier for others to understand the problem you are having.

Thank you.

3

I think using numba.np.unsafe.ndarray.to_fixed_tuple() it should be possible to generate a tuple dynamically from an array as long as the length is a compile time constant (eg. len(a)).

Also @stuartarchibald showed me a method on Jan 13 in the Gitter chat (I don’t know how to generate a link)

Also have a look here: https://github.com/numba/numba/issues/4265

1 Like
4

Thank you for your hint.

>>> import numba
>>> import numpy as np
>>> from numba.np.unsafe.ndarray import to_fixed_tuple
>>> arr = np.arange(8).reshape(2,2,2)
>>> @numba.njit
... def transpose(a):
...     axes = np.arange(a.ndim)
...     axes[: 2] = [1, 0]
...     axes = to_fixed_tuple(axes, a.ndim)
...     return np.transpose(a, axes)

# this works now
>>> transpose(a)

If I try to use this snippet within the numba core package, it fails due to non constant length, see MR [WIP] Add support for np.rot90 by braniii · Pull Request #6822 · numba/numba · GitHub for the source code. Is there a way to avoid that?

$ python -m numba.runtests numba.tests.test_np_functions.TestNPFunctions.test_rot90_basic
E
======================================================================
ERROR: test_rot90_basic (numba.tests.test_np_functions.TestNPFunctions)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/braniii/Documents/Coding/numba/numba/tests/test_np_functions.py", line 2336, in test_rot90_basic
    got = cfunc(a, k)
  File "/home/braniii/Documents/Coding/numba/numba/core/dispatcher.py", line 421, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/braniii/Documents/Coding/numba/numba/core/dispatcher.py", line 362, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function rot90 at 0x7f2b7f4380d0>) found for signature:
 
 >>> rot90(array(int64, 2d, C), int64)
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'numpy_rot90': File: numba/np/arrayobj.py: Line 1566.
    With argument(s): '(array(int64, 2d, C), int64)':
   Rejected as the implementation raised a specific error:
     TypingError: Failed in nopython mode pipeline (step: nopython frontend)
   No implementation of function Function(<intrinsic to_fixed_tuple>) found for signature:
    
    >>> <unknown function>(array(int64, 1d, C), int64)
    
   There are 2 candidate implementations:
         - Of which 2 did not match due to:
         Intrinsic in function 'to_fixed_tuple': File: numba/np/unsafe/ndarray.py: Line 41.
           With argument(s): '(array(int64, 1d, C), int64)':
          Rejected as the implementation raised a specific error:
            RequireLiteralValue: *length* argument must be a constant
     raised from /home/braniii/Documents/Coding/numba/numba/np/unsafe/ndarray.py:52
   
   During: resolving callee type: Function(<intrinsic to_fixed_tuple>)
   During: typing of call at /home/braniii/Documents/Coding/numba/numba/np/arrayobj.py (1589)
   
   
   File "numba/np/arrayobj.py", line 1589:
       def impl(arr, k=1):
           <source elided>
           axes_list[:2] = [1, 0]
           axes_list = to_fixed_tuple(axes_list, arr.ndim)
           ^

  raised from /home/braniii/Documents/Coding/numba/numba/core/typeinfer.py:1071

During: resolving callee type: Function(<function rot90 at 0x7f2b7f4380d0>)
During: typing of call at /home/braniii/Documents/Coding/numba/numba/tests/test_np_functions.py (133)


File "numba/tests/test_np_functions.py", line 133:
def rot90(a, k=1):
    return np.rot90(a, k)

Edit: Sorry I just understood what you meant with Gitter. So I guess I need to use @overload instead?

5

so you cannot fix the number of dimensions (length of tuple) at compile time. can you narrow down the options? For example, if it can only be 4, 5, 6 or 7?

6

For me it would be fine to limit the dimension, but I guess it will not be merged if it can be applied only to low dimensional arrays.

Actually I do not understand why it works in the MWE. Does it get compiled for each dimension separately? And why is the behaviour different when implementing it in numba directly? I need to admit, that I am new to numba and still not understand the way it works.

7

Does it get compiled for each dimension separately? And why is the behaviour different when implementing it in numba directly? I need to admit, that I am new to numba and still not understand the way it works.

once you understand why that happens, you’ll understand the most important part of numba from the point of view of the user.
Maybe this can help a bit: Polymorphic dispatching — Numba 0+untagged.4124.gd4460fe.dirty documentation

when you pass axes as an argument numba is able to look at your tuple and compile the function for that specific length. you can verify that by running print(transpose1.signatures). At first it will be empty, and then for each execution with a new length of tuple you will see a new signature coming up.
When you try to create the tuple inside the compiled code this fails, because it cannot be compiled before knowing the length of the tuple.
The length of the tuple you are trying to create could be inferred from the ndim of the a array, but the connection is too complex for the compiler to figure it out. Conceptually it’s possible, but you will have to do part of the work to make it obvious to the compiler. Numba will compile one version of the function for each ndim that you pass.

There are several ways to do it, one of them is using generated_jit. Something like,

@numba.generated_jit
def my_transpose(a):
    
    ndim = a.ndim
    
    def impl_transpose(a):
        axes = np.arange(ndim)
        return to_fixed_tuple(axes, ndim)
    
    return impl_transpose

by making ndim a closure variable the compiler knows that it’s a constant.

Hope this helps

Luk

1 Like
8

Think it’s this method: https://gitter.im/numba/numba?at=5ffef04791e9b71badd35eaf

1 Like
9

This fixes the problem. Thanks all, this is a great community here :slight_smile:

10

And thx for the short overview. I’ve started to read the doc. But I guess it needs some time and experience to understand what is written there.