在 numba njit 中创建具有指定顺序的 numpy 数组副本

Question

我想创建一个传递给 numba njitted 函数的数组副本，并将该数组的顺序固定为 Fortran 顺序。总之，我想实现这个：

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

但在 njitted 函数中。

到目前为止我尝试过的：

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)

最新的尝试可以工作，但似乎没有任何方法可以检查 numba 中现有数组的顺序（不支持参数“flags”）。

更新：

import numpy as np
from numba import njit, generated_jit

@generated_jit(nopython=True)
def fortran_copy(a):
    if a.is_f_contig:
        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

更新 2：经过一些测试，我发现检查 a.is_f_contig（或 a.layer == 'F'）标志是不够的。有时该函数不会 return 副本。我不确定这里有什么问题。

Answer 1

您可以使用 foo.nopython_signatures 检查函数的 input/output 类型。数组的顺序是数组类型的一部分，并按原样报告。请注意，在编译函数中在运行时获取类型是一种称为内省的功能，并且通常几乎不支持原生 statically-typed 语言，如 C 或 C++（与 dynamically-typed 语言，如 Python 相对） .这就是 Numba 中的 AFAIK 情况（尽管它可能会在不久的将来发生变化，因为最近实施了 isinstance）。解决这个问题的方法是使用 @generated-jit 函数。这是此功能的 documentation。在这样的函数中，您可以使用 b.flags['F_CONTIGUOUS'] 检查数组的输入类型，然后 return return FORTRAN 数组的最佳函数（即 returning np.asfortranarray(a) 或简单地 np.copy(a)).

在 numba njit 中创建具有指定顺序的 numpy 数组副本

Create copy of numpy array with specified order inside numba njit

python

numpy

numba