Create copy of numpy array with specified order inside numba njit-CodePudding

I would like to create a copy of an array that was passed into numba njitted function and fix an order of this array to a Fortran one. In short, I would like to achieve this:

np.copy(a, order='F')

But inside njitted function.

What I have tried so far:

import numpy as np
import numba as nb
from numba import njit

a = np.random.randn(10,5)
@njit
def foo(a):
    b = np.require(a=a, requirements='F_CONTIGUOUS')
    b = np.asarray(a=a, order='F')
    b = np.copy(a, order='F')
    if nb.typeof(a).layout != 'F':
        b = np.asfortranarray(a)
    else:
        b = np.copy(a)
    return b
foo(a)

The latest try could work, but it does not seem that there is any way to check the order of an existing array in numba (parameter "flags" is not supported).

UPDATE:

import numpy as np
from numba import njit, generated_jit

@generated_jit(nopython=True)
def fortran_copy(a):
    if a.layout == 'F':
        return lambda a: np.copy(a)
    else:
        return lambda a: np.asfortranarray(a)

a = np.random.randn(10,5,1)
b = np.asfortranarray(a)
c = fortran_copy(a)
assert not np.may_share_memory(a,c)
assert not c.base is a
c = fortran_copy(b)
assert not np.may_share_memory(b,c)
assert not c.base is b

CodePudding user response：

You can check the input/output types of a function with foo.nopython_signatures. The ordering of an array is part of an array type and is reported as such. Note that getting the type at runtime in a compiled function is a feature called introspection and is often barely supported in native statically-typed language like C or C (as opposed to dynamically-typed languages like Python). This is AFAIK the case in Numba (though it may change in a near future since isinstance has been recently implemented). The solution to address this problem is to use @generated-jit function. Here is the documentation for this feature. In such function, you can check the input type of an array with b.flags['F_CONTIGUOUS'] and then return the best function to return a FORTRAN array (ie. returning either np.asfortranarray(a) or simply np.copy(a)).