Python 函数：return 一个没有额外内存分配的数组

Question

假设我想创建一个将输入向量乘以输入矩阵的函数：

def MatMul(A,b):
    return A.dot(b)

现在，我执行以下代码：

import numpy as np
A=np.array([[1,2,3],[4,5,6],[7,8,9]],dtype='float64')
b=np.array([4,5,6],dtype='float64')
c=np.zeros(3,dtype='float64')
c=MatMul(A,b)

MatMul函数内部会不会有额外的数组分配？我知道 A 和 b 将通过引用传递。请注意，我已经预先分配了数组 c。

一般来说，如何避免在像这样的简单函数中进行不必要的预分配？比如说，我想执行几个数学运算：

def Rank1Update(A,b,alpha):
    c=A.dot(b)
    c+=alpha*c.dot(c)*c
    return c

我可以在 1 行中包含许多数学函数，但代码很快变得不可读。

我熟悉 C 风格的编程，为了避免不必要的内存分配，可以通过引用传递 A、b 和 c 并更新 c returns void 的函数里面。我可以在 python 中做同样的事情，但为了方便和代码可读性，我想使用 return，

谢谢，

米哈伊尔

Answer 1

dot 不知道也不关心 c 变量或该变量已经包含引用的数组。它将创建一个新数组，= 会将 c 变量绑定到该新数组，让内存管理系统清理旧数组。

如果您希望 dot 将输出存储到现有数组中，您需要告诉它这样做：

A.dot(b, out=c)

Answer 2

Numpy 在其大多数生成新数组的函数中都有一个参数来存储生成的数组。我在下面的代码中使用了该参数的名称版本 out，但您可以省略名称。您必须确保 out 数组具有正确的形状和数据类型。这个参数的目的正是你想要的——避免额外的内存分配。这样也可以加快代码速度。

import numpy as np
A = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype='float64')
b = np.array([4, 5, 6], dtype='float64')
c = np.zeros(3, dtype='float64')

A.dot(b, out=c)

documentation for dot() 中提到了该参数。如果你想改变函数的定义 MatMul,

def MatMul(A, b, c=None):
    return A.dot(b, out=c)

并将调用更改为

MatMul(A, b, c)

Python 函数：return 一个没有额外内存分配的数组

Python function: return an array without additional memory allocations

python

numpy

function

pass-by-reference