Numpy 标量上的二进制操作自动向上转换为 float64答案

【问题标题】：Binary operations on Numpy scalars automatically up-casts to float64Numpy 标量上的二进制操作自动向上转换为 float64
【发布时间】：2014-11-13 10:07:59
【问题描述】：

我想在 np.float32 和内置 Python int 和 float 之间进行二进制操作（如加法和乘法），并获得 np.float32 作为返回类型。但是，它会自动向上转换为 np.float64。

示例代码：

>>> a = np.float32(5)
>>> a.dtype
dtype('float32')
>>> b = a + 2
>>> b.dtype
dtype('float64')

如果我使用 np.float128 执行此操作，b 也会变成 np.float128。这很好，因为它因此保持了精度。但是，在我的示例中，不需要向上转换到 np.float64 来保持精度，但它仍然会发生。如果我将 2.0（一个 Python 浮点数（64 位））添加到 a 而不是 2，那么强制转换是有意义的。但即使在这里，我也不要它。

所以我的问题是：在将二元运算符应用于 np.float32 和内置 Python int/float 时，如何更改已完成的转换？或者，将单精度作为所有计算的标准而不是双精度，也可以算作一种解决方案，因为我不需要双精度。别人问过这个问题，好像没有找到解决办法。

我知道 numpy 数组和 dtypes。在这里，我得到了想要的行为，因为数组始终保留其 dtype。然而，当我对数组的单个元素进行操作时，我得到了不需要的行为。我对解决方案有一个模糊的想法，涉及子类化 np.ndarray（或 np.float32）并更改 __array_priority__ 的值。到目前为止，我还不能让它工作。

我为什么要关心？我正在尝试使用 Numba 编写一个 n-body 代码。这就是为什么我不能简单地对整个数组进行操作。将所有 np.float64 更改为 np.float32 可使速度提高约 2 倍，这很重要。 np.float64-casting 行为完全破坏了这种加速，因为我的 np.float32 数组上的所有操作都以 64 精度完成，然后向下转换为 32 精度。

【问题讨论】：

标签： python numpy numba single-precision

【解决方案1】：

我不确定 NumPy 的行为，或者您究竟是如何尝试使用 Numba，但明确说明 Numba 类型可能会有所帮助。例如，如果您执行以下操作：

@jit
def foo(a):
    return a[0] + 2;

a = np.array([3.3], dtype='f4')
foo(a)

a[0] 中的 float32 值在添加操作之前被提升为 float64（如果您不介意深入研究 llvm IR，您可以通过使用 numba 命令运行代码并使用 - -dump-llvm 或 --dump-optimized 标志：numba --dump-optimized numba_test.py）。但是，通过指定函数签名，包括返回类型为 float32：

@jit('f4(f4[:]'))
def foo(a):
    return a[0] + 2;

a[0] 中的值不会提升为 float64，尽管结果被强制转换为 float64，因此当函数返回 Python 时它可以转换为 Python 浮点对象。

如果你可以预先分配一个数组来保存结果，你可以这样做：

@jit
def foo():
    a = np.arange(1000000, dtype='f4')
    result = np.zeros(1000000, dtype='f4')
    for i in range(a.size):
        result[0] = a[0] + 2

即使您自己进行循环，编译代码的性能也应该与 NumPy ufunc 相当，并且不应发生到 float64 的强制转换（同样，这可以通过查看 Numba 生成的 llvm IR 来验证)。

【讨论】：