Numba JIT模式下从现有NumPy数组创建新数组的正确姿势

本文探讨Numba JIT编译模式下，直接使用`np.array(existing_array)`从现有NumPy数组创建新数组时遇到的`TypingError`。文章将澄清此问题与Numba字典无关，而是`np.array()`构造函数的特定限制，并提供通过解包操作符`*`或适当的构造方法来解决此问题的专业指导，确保代码在Numba环境中高效运行。

在高性能计算领域，Numba通过即时编译（JIT）技术显著提升Python代码的执行效率，尤其在处理NumPy数组时表现出色。然而，在使用Numba的nopython模式时，开发者可能会遇到一些特定的类型推断和函数实现限制。其中一个常见的困惑是，当尝试使用np.array()构造函数从一个已存在的NumPy数组创建另一个NumPy数组时，Numba会抛出TypingError。

Numba中np.array()构造函数的限制解析

初看之下，这个错误可能让人误以为是Numba对字典值类型的特殊处理，但实际上，它与Numba如何处理np.array()构造函数有关。Numba的nopython模式需要所有操作都有明确的类型签名。当您尝试将一个NumPy数组作为参数直接传递给np.array()时，例如np.array(a)，其中a本身就是一个np.ndarray，Numba会报告找不到匹配的函数实现。

考虑以下示例，它展示了在Numba JIT编译函数中直接使用np.array(a)引发的错误：

import numpy as np
import numba as nb

@nb.njit
def problematic_foo(a):
    # 尝试从现有NumPy数组 'a' 创建一个新的NumPy数组 'x'
    x = np.array(a) # 此处会引发TypingError
    return x

# 示例调用
a_data = np.array([1, 2, 3], dtype=np.int64)
try:
    problematic_foo(a_data)
except Exception as e:
    print(f"发生错误: {e}")

运行上述代码，您会看到一个TypingError，其中关键信息是： No implementation of function Function(<built-in function array>) found for signature: >>> array(array(int64, 1d, C)) 这明确指出Numba在处理np.array(array(int64, 1d, C))这种签名时遇到了障碍。Numba的np.array()实现通常期望接收一个可迭代对象（如列表、元组），其中包含可以转换为标量类型的数据，而不是另一个完整的NumPy数组对象。

根本原因分析

Numba在nopython模式下工作时，会对代码进行静态类型推断和编译。它维护了一套其支持的函数和操作的内部实现。对于np.array()，Numba的内部实现主要针对以下几种情况：

1. 从Python列表或元组创建数组：np.array([1, 2, 3])
2. 从标量值创建数组：np.array(5)
3. 指定数据类型或维度创建空数组：np.empty(shape, dtype)

然而，Numba当前版本并未提供一个直接的、优化过的np.array(existing_np_array)实现，即从一个NumPy数组对象本身构造一个新的NumPy数组。它将existing_np_array视为一个单一的、不可迭代的“对象”来处理，而不是将其内部元素提取出来进行构造。

怪兽AI数字人

数字人短视频创作，数字人直播，实时驱动数字人

44

查看详情

解决方案：正确创建NumPy数组

要解决这个问题，我们需要确保传递给np.array()的是Numba能够理解和处理的可迭代对象，例如一个包含原始数组元素的Python列表。最简洁且推荐的方法是使用Python的解包操作符*将现有NumPy数组的元素解包到一个列表中，然后再将该列表传递给np.array()。

以下是修正后的代码示例：

import numpy as np
import numba as nb

@nb.njit
def correct_foo(a, b, c):
    # 假设 'a' 是一个NumPy数组
    # 使用解包操作符 '*' 将 'a' 的元素解包成一个列表
    # 然后 np.array() 可以从这个列表中创建新数组
    x = np.array([*a])

    # 验证这个操作在Numba字典中也适用
    d = {}
    d[(1, 2, 3)] = x # 现在 'x' 是一个有效的NumPy数组，可以作为字典值
    return d

# 示例调用
a_data = np.array([1, 2], dtype=np.int64)
b_data = np.array([3, 4], dtype=np.int64)
c_data = 5 # 假设 c 是一个标量，虽然在这个例子中未使用
result_dict = correct_foo(a_data, b_data, c_data)
print(result_dict)
# 预期输出: {(1, 2, 3): array([1, 2])}

在这个correct_foo函数中，np.array([*a])的工作原理是：

1. *a将NumPy数组a的元素解包。对于一维数组，这会产生一系列独立的元素。
2. [*a]将这些独立的元素收集到一个Python列表中。
3. np.array(...)现在接收到一个标准的Python列表，Numba对此有明确的实现，可以成功地从该列表创建新的NumPy数组。

注意事项与性能考量

1. 理解Numba的类型推断: Numba的强大之处在于其静态类型推断。当遇到TypingError时，通常意味着您正在尝试执行一个Numba没有明确实现或不支持的操作签名。
2. 区分复制与创建:
   • 如果您的目标仅仅是创建一个现有NumPy数组的副本，更高效的方法是使用a.copy()或np.copy(a)。这些方法通常在Numba中得到良好支持，且避免了创建中间Python列表的开销。
   • np.array([*a])虽然解决了问题，但在处理非常大的数组时，创建中间Python列表可能会引入额外的内存开销和一定的性能损耗。因此，在性能敏感的场景下，应优先考虑a.copy()。
3. 字典与数组: Numba字典可以很好地存储NumPy数组作为其值，前提是这些数组本身是Numba能够正确处理的类型。本教程澄清了问题不在于字典本身，而在于数组的构造方式。

总结

在Numba的nopython模式下，直接使用np.array(existing_np_array)构造新数组会导致TypingError，因为它没有匹配的函数签名实现。正确的做法是利用Python的解包操作符*将现有数组的元素转换为一个列表，例如np.array([*existing_np_array])。然而，如果仅仅是为了复制数组，existing_np_array.copy()或np.copy(existing_np_array)是更直接和高效的选择。理解Numba的类型系统和其对NumPy操作的特定支持是编写高效JIT编译代码的关键。

以上就是Numba JIT模式下从现有NumPy数组创建新数组的正确姿势的详细内容，更多请关注php中文网其它相关文章！