Numba njit 中 np.array() 使用误区与解决方法

在Numba JIT编译函数中，直接使用`np.array(existing_numpy_array)`构造新数组会引发`TypingError`，因为Numba的`np.array`实现不支持以现有NumPy数组作为其唯一参数来重新创建数组。该问题与Numba字典无关，而是NumPy数组构造方式的限制。解决方案是使用解包操作符，如`np.array([*existing_numpy_array])`，以正确传递数组元素。

在高性能计算领域，Numba通过即时编译（JIT）技术显著提升Python代码，特别是涉及NumPy数组操作时的执行效率。然而，在使用Numba的@njit装饰器时，开发者可能会遇到一些与NumPy数组构造相关的细微差别，这些差别可能导致TypingError。本文将深入探讨一个常见的误区：在Numba JIT编译环境中，如何正确地从一个现有NumPy数组创建另一个NumPy数组，以及避免由此产生的TypingError。

理解Numba中np.array()构造器的行为

当我们在Numba JIT编译的函数中操作NumPy数组时，需要特别注意np.array()构造函数的行为。一个常见的误解是，可以像在纯Python/NumPy环境中那样，直接将一个现有的NumPy数组作为参数传递给np.array()来创建一个新的数组。例如，以下代码在纯Python环境中通常是可行的：

import numpy as np
a = np.array([1, 2, 3])
b = np.array(a) # b是a的一个副本

然而，在Numba的@njit编译模式下，这种做法会引发TypingError。

示例：导致TypingError的代码

考虑以下Numba JIT编译函数中的代码片段：

import numpy as np
import numba as nb

@nb.njit
def create_array_from_existing(arr_in):
    # 尝试直接从现有NumPy数组arr_in创建新数组
    new_arr = np.array(arr_in)
    return new_arr

# 调用函数
a = np.array([1, 2])
# create_array_from_existing(a) # 这行代码将导致TypingError

执行上述代码，Numba会抛出类似于以下的TypingError：

TypingError: No implementation of function Function() found for signature:
>>> array(array(int64, 1d, C))

错误信息清晰地指出，Numba无法找到一个np.array()的实现，其签名接受一个array(int64, 1d, C)类型的参数。这意味着Numba的np.array()构造器在njit模式下，并不支持直接以一个NumPy数组对象作为其唯一的输入参数来重新创建数组。它期望的输入通常是一个Python序列（如列表、元组）或可迭代对象，其中包含可以转换为数组元素的标量值。

进一步的错误信息可能会显示：array(int64, 1d, C) not allowed in a homogeneous sequence。这表明Numba试图将整个NumPy数组对象arr_in视为一个序列中的单个元素，而不是将其内部元素提取出来构成新数组。

解决方案：使用解包操作符

为了在Numba JIT编译函数中正确地从一个现有NumPy数组的元素创建新数组，我们需要明确地将原数组的元素解包（unpack）到一个Python序列中，然后将该序列传递给np.array()。Python的解包操作符*在这里发挥了关键作用。

import numpy as np
import numba as nb

@nb.njit
def create_array_from_existing_fixed(arr_in):
    # 使用解包操作符将arr_in的元素传递给np.array()
    new_arr = np.array([*arr_in])
    return new_arr

# 调用函数
a = np.array([1, 2])
result = create_array_from_existing_fixed(a)
print(result)
print(type(result))

输出：

[1 2]

通过np.array([*arr_in])，我们将arr_in中的所有元素解包到一个临时的Python列表中，然后np.array()能够正确地从这个列表构造一个新的NumPy数组。这种方法符合Numba对np.array()构造器参数类型的期望。

注意事项与最佳实践

1. 区分复制与重新创建：

   • 如果你仅仅需要一个现有NumPy数组的副本，最直接且推荐的方法是使用arr_in.copy()。这通常更高效，并且明确表达了意图。

     @nb.njit
     def copy_array(arr_in):
     return arr_in.copy()

   • np.array([*arr_in])实际上创建了一个新的数组，其元素与原数组相同。如果需要更改数据类型或数组布局，这种方法结合dtype参数会很有用。

2. Numba字典与NumPy数组值： 最初的问题中提到NumPy数组作为字典值。实际上，Numba字典可以很好地存储NumPy数组作为其值，只要这些数组的类型在编译时是同质的（homogeneous）。本例中的TypingError与字典本身无关，而是np.array()构造函数的特定行为导致的。一旦NumPy数组被正确创建，将其赋值给Numba字典的值是完全支持的。

   @nb.njit
   def foo_with_dict(a):
       d = nb.typed.Dict.empty(
           key_type=nb.types.int64,
           value_type=nb.typeof(np.array([0])) # 定义字典值的类型
       )
       d[1] = np.array([*a]) # 正确创建并赋值
       return d
   
   a = np.array([1, 2])
   my_dict = foo_with_dict(a)
   print(my_dict)

3. 理解Numba的类型推断： Numba在编译时会进行严格的类型推断。当遇到np.array(arr_in)时，它推断arr_in是一个NumPy数组类型，并尝试寻找一个接受NumPy数组作为参数的np.array重载。由于Numba没有这样的直接重载来“复制”一个数组，因此会报错。而[*arr_in]会产生一个临时的Python列表，Numba能够识别并处理np.array(list_of_scalars)的签名。

总结

在Numba的@njit编译环境中，直接使用np.array(existing_numpy_array)来创建新数组会导致TypingError。核心原因在于Numba对np.array()构造器参数类型的严格要求，它不直接支持以一个完整的NumPy数组对象作为输入来重新创建数组。正确的做法是利用Python的解包操作符*，将现有NumPy数组的元素解包成一个列表，然后将该列表传递给np.array()，即np.array([*existing_numpy_array])。此外，如果仅仅需要数组的副本，使用existing_numpy_array.copy()是更简洁高效的选择。理解这些细节对于在Numba中编写高效且无误的NumPy代码至关重要。