如何方便地在 numpy fortran contiguos 数组上使用操作？

Question

一些像np.matmul(a, b)这样的numpy函数对矩阵堆栈有方便的行为。

手册指出：

如果任一参数为 N-D，N > 2，则将其视为驻留在最后两个索引中的矩阵堆栈并相应地广播。

因此，对于a.shape = (10 , 2, 4) 和b.shape(10, 4, 2)，语句a @ b 是有意义的并且将具有(10, 2, 2) 的形状

但是，我来自线性代数世界，在那里我习惯了 Fortran 连续数组布局。

表示为 Fortran 连续数组的相同 a 将具有形状 (4, 2, 10) 和类似的 b.shape = (2, 4, 10)。

要像以前一样执行a @ b，我必须调用 (a.T @ b.T).T.

更糟糕的是，假设您天真地创建了相同的 Fortran 连续数组 a，并考虑了 matmul 的行为，因此它的形状为 (10, 4, 2)。 然后a.strides = (8, 80, 320) 在“堆栈”索引中具有最小的步幅，实际上应该具有最大的步幅。

这真的是要走的路还是我错过了什么？

Answer 1

numpy 的矩阵乘法与数组的内部布局无关。例如，这里有两个 C 有序数组：

>>> import numpy as np
>>> a = np.random.rand(10, 2, 4)
>>> b = np.random.rand(10, 4, 2)
>>> print('a', a.shape, a.strides)
>>> print('b', b.shape, b.strides)
a (10, 2, 4) (64, 32, 8)
b (10, 4, 2) (64, 16, 8)

以下是 Fortran 顺序的等效数组：

>>> af = np.asfortranarray(a)
>>> bf = np.asfortranarray(b)
>>> print('af', af.shape, af.strides)
>>> print('bf', bf.shape, bf.strides)
af (10, 2, 4) (8, 80, 160)
bf (10, 4, 2) (8, 80, 320)

Numpy 将等价的数组视为等价的，不管它们的内部布局如何：

>>> np.allclose(a, af) and np.allclose(b, bf)
True

矩阵乘法的结果不依赖于内部布局：

>>> np.allclose(a @ b, af @ bf)
True

如果你愿意，你甚至可以混合布局：

>>> np.allclose(a @ bf, af @ b)
True

简而言之，在 numpy 中使用 Fortran 有序数组最方便的方法是不用担心内部数组布局：形状才是最重要的。

如果您的数组形状与numpy matmul API 所期望的不同，您最好的办法是重塑数组，例如使用a.transpose(2, 0, 1) @ b.transpose(2, 0, 1) 或类似的，具体取决于您的用例，但不要'不用担心：对于 C 或 Fortran 连续数组，此操作仅调整数组视图周围的元数据，不会导致底层数据缓冲区被复制或重新排序。

Answer 2

虽然numpy 可以处理各种布局，但许多细节的设计都考虑到了“C”布局。很好的例子是嵌套列表如何转换为数组，以及 numpy 操作在 matmul 案例中批量处理多余维度的方式。

根据经验，numpy 的结果不依赖于数组布局（FORTRAN，C，非连续）是正确的；然而，速度确实如此，而且非常重要：

rng = np.random.default_rng()
a = rng.random((100,111,200))
b = rng.random((111,77,200))
af = np.array(a,order="F")
bf = np.array(b,order="F")

np.allclose((b.T@a.T).T,(bf.T@af.T).T)
# True
timeit(lambda:(b.T@a.T).T,number=10)
# 5.972857117187232
timeit(lambda:(bf.T@af.T).T,number=10)
# 0.1994628761895001

事实上，有时候非懒惰的转置是完全值得的，即将你的数据复制到最好的布局中：

timeit(lambda:(np.array(b.T,order="C")@np.array(a.T,order="C")).T,number=10)
# 0.3931349152699113

我的建议：如果您想要速度和方便，最好使用“C”布局，它不会花费很长时间来适应并且可以为您省去很多潜在的麻烦。