我不喜欢 NumPy

我不喜欢 NumPy

2025年5月

评论在 lemmmy, substack。 人们常说，除非你先爱上某个人，否则你无法真正恨他。我不知道这作为一个普遍原则是否正确，但它确实描述了我与 NumPy 的关系。 顺便说一下，NumPy 是一些在 Python 中对数组进行计算的软件。它非常流行，并且对所有流行的机器学习库（如 PyTorch）产生了巨大的影响。这些库也存在我下面讨论的大部分问题，但我为了具体起见，将只讨论 NumPy。

NumPy 让简单的事情变得简单。比如 A 是一个 5×5 矩阵，x 是一个长度为 5 的向量，你想找到一个向量 y 使得 Ay=x。在 NumPy 中，这会是：

y = np.linalg.solve(A, x)

多么优雅！多么清晰！

但假设情况稍微复杂 一点。假设 A 是一个由 100 个 5×5 矩阵堆叠而成的矩阵，表示为一个 100×5×5 的数组。并且假设 x 是一个由 100 个长度为 5 的向量堆叠而成的向量，表示为一个 100×5 的数组。并且假设你想求解 Aᵢyᵢ=xᵢ，其中 1≤i≤100。

如果你能使用循环，这会很容易：

y = np.empty_like(x)
for i in range(100):
  y[i,:] = np.linalg.solve(A[i,:,:], x[i,:])

但是你不能使用循环。在某种程度上，这是因为 Python 中的循环速度很慢所带来的限制。但如今，一切都在 GPU 上运行，如果你有很大的数组，你可能不想在任何语言中使用循环。为了让所有这些晶体管都运行起来，你需要调用特殊的 GPU 函数，这些函数会将数组分割成许多小块并并行处理它们。

好消息是 NumPy 知道这些特殊例程（至少如果你使用 JAX 或 CuPy），并且如果你正确地调用 np.linalg.solve，它将使用它们。

坏消息是没人知道如何做到这一点。

不相信我？好吧，那么下面哪个是对的？

y = linalg.solve(A,x)
y = linalg.solve(A,x,axis=0)
y = linalg.solve(A,x,axes=[[1,2],1])
y = linalg.solve(A.T, x.T)
y = linalg.solve(A.T, x).T
y = linalg.solve(A, x[None,:,:])
y = linalg.solve(A,x[:,:,None])
y = linalg.solve(A,x[:,:,None])[:,:,0]
y = linalg.solve(A[:,:,:,None],x[:,None,None,:])
y = linalg.solve(A.transpose([1,2,0]),x[:,:,None]).T

没人知道。并且让我向你展示另一件事。这是 documentation：

np.linalg.solve

阅读它。沉思它。现在，请注意：你 仍然 不知道如何一次性地求解所有 i 的 Aᵢyᵢ=xᵢ。这甚至可能吗？当我说它可能的时候，我撒谎了吗？

据我所知，人们实际上所做的是尝试随机变体，直到有一个似乎可以工作为止。

为什么 NumPy 不好

NumPy 的全部意义在于将运算应用于数组。当数组具有 2 个或更少的维度时，一切都很好。但是如果你正在做一些稍微复杂的事情，你不可避免地会发现自己想要将某些操作应用于数组 A 的某些维度、数组 B 的某些其他维度，以及数组 C 的一些 其他 维度。而 NumPy 没有关于如何表达这一点的理论。

让我告诉你我的意思。假设：

并且假设对于每个 k 和 n，你想计算 L 和 M 维度上的平均值。也就是说，你想要 Dkn = 1/(LM) × ∑lm Aklm Bln Ckm. 为了做到这一点，你有两个选择。第一个是使用怪异的维度对齐技巧：

D = np.mean(
    np.mean(
      A[:,:,:,None] *
      B[None,:,None,:] *
      C[:,None,:,None],
    axis=1),
  axis=1)

你会问，这是什么鬼？为什么到处都是 None？好吧，在 NumPy 中对数组进行索引时，你可以写入 None 以插入一个新维度。 A 是 K×L×M，但 A[:,:,:,None] 是 K×L×M×1。同样，B[None,:,None,:] 是 1×L×1×N，C[:,None,:,None] 是 K×1×M×1。当你将它们相乘时，NumPy 会“广播”所有大小为 1 的维度，以给出一个 K×L×M×N 数组。然后，np.mean 调用会对 L 和 M 维度求平均值。

我认为这很糟糕。我使用 NumPy 已经很多年了，但我仍然觉得不犯错地编写这样的代码是不可能的。

它也几乎不可能阅读。为了证明这一点，我刚刚抛了一枚硬币，当且仅当硬币是反面时，我在上面引入了一个错误。有错误吗？你 确定 吗？没人知道。

你的第二个选择是拼命地想变得聪明。生命短暂而宝贵，但如果你花很多时间阅读 NumPy 文档，你最终可能会意识到有一个名为 np.tensordot 的函数，并且可以使其完成大部分工作：

D = (1/L) * np.mean(
        np.tensordot(A, B, axes=[1,0]) *
        C[:,:,None],
      axis=1)

这是正确的。（我保证。）但为什么它会起作用？ np.tensordot 到底在做什么？如果你在其他上下文中看到该代码，你会对正在发生的事情有丝毫的了解吗？

如果我可以使用循环，我会这样做：

D = np.zeros((K,N)) 
for k in range(K): 
  for n in range(N): 
    a = A[k,:,:] 
    b = B[:,n] 
    c = C[k,:] 
    assert a.shape == (L,M) 
    assert b.shape == (L,) 
    assert c.shape == (M,) 
    D[k,n] = np.mean(a * b[:,None] * c[None,:])

写了太多 NumPy 的人可能会觉得这很笨拙。我怀疑这有点斯德哥尔摩综合症。但我们肯定可以同意它是 清晰的。

在实践中，情况通常更糟。假设 A 的形状是 M×K×L 而不是 K×L×M。使用循环，没什么大不了的。但是 NumPy 要求你编写像 A.transpose([1,2,0]) 这样的怪物。或者应该是 A.transpose([2,0,1]) 吗？这些产生什么形状？没人知道。

循环更好。

好吧，我撒谎了

还有第三个选择：

D = 1/(L*M) * np.einsum('klm,ln,km->kn', A, B, C)

如果你以前从未见过 Einstein 求和，那可能看起来很可怕。但请记住，我们的目标是找到 Dkn = 1/(LM) × ∑lm Aklm Bln Ckm. 上面代码中的字符串基本上给出了每个三个输入（klm,ln,km）中的索引标签，以及输出的目标索引（->kn）。然后，np.einsum 将输入的对应元素相乘，并对输出中不存在的所有索引求和。

就我个人而言，我认为 np.einsum 是 NumPy 中为数不多的真正好的部分之一。这些字符串有点繁琐，但它们是值得的，因为整个函数易于理解，完全明确，并且非常通用和强大。

除了，np.einsum 如何实现所有这些？它使用索引。或者，更准确地说，它引入了一种基于索引的微型 特定领域语言。它没有受到 NumPy 的设计缺陷的影响，因为它拒绝遵守 NumPy 的正常规则。

但是 np.einsum 只做一些事情。（Einops 做的更多一些。）如果你想将一些其他函数应用于某些数组的各个维度呢？没有 np.linalg.einsolve。并且如果你创建自己的函数，那么 肯定 没有该函数的“Einstein”版本。

我认为 np.einsum 的优点表明 NumPy 走对了一些地方。

幕间休息

这是一幅 painting，它给人的感觉与我们的主题类似。

NumPy 哪里出错了？

这是我对数组语言的期望。我并不特别在意语法，但如果能做到以下几点就太好了：

  1. 当你想做某事时，如何做这件事是“显而易见”的。
  2. 当你阅读一些代码时，它所做的事情是“显而易见”的。

那不是很好吗？我认为 NumPy 没有实现这些目标，因为它最初的罪过是：它取消了索引并用 broadcasting 替换了它们。而 broadcasting 无法填补索引的空白。

我不喜欢 NumPy 的 broadcasting

NumPy 的核心技巧是 broadcasting。看这段代码：

A = np.array([[1,2],[3,4],[5,6]])
B = np.array([10,20])
C = A * B
print(C)

这输出：

[[ 10 40]
 [ 30 80]
 [ 50 120]]

这里，A 是一个 3×2 数组，B 是一个长度为 2 的数组。当你将它们相乘时，B 被“广播”到 A 的形状，这意味着 A 的第一列与 B[0]=10 相乘，第二列与 B[1]=20 相乘。

在简单的情况下，这似乎不错。但是我不喜欢它。一个原因是，正如我们上面看到的，你经常需要对维度做一些糟糕的事情才能使它们对齐。

另一个原因是它不明确或不清晰。有时 A*B 逐元素相乘，有时它会做更复杂的事情。因此，每次你看到 A*B 时，你都必须弄清楚 broadcasting conventions 中的哪种情况被触发了。

但是 broadcasting 的真正问题在于它如何感染其他一切。我将在下面解释。

我不喜欢 NumPy 的 indexing

这是一个谜题。看这段代码：

A = np.ones((10,20,30,40))
i = np.array([1,2,3])
j = np.array([[0],[1]])
B = A[:,i,j,:]

B 的形状是什么？

事实证明答案是 10×2×3×40。那是因为 i 和 j 索引被广播到 2×3 的形状，然后是这个那个，嘟囔嘟囔嘟囔。尝试说服自己这是有道理的。

完成了？好的，现在试试这些：

C = A[:,:,i,j]
D = A[:,i,:,j]
E = A[:,1:4,j,:]

这些的形状是什么？

是的，这就是发生的事情。如果你不相信我，请尝试一下！我理解 NumPy 为什么这样做，因为我吸收了 this 5000 word document，它解释了 NumPy 索引的工作方式。但是我想找回那段时间。

为了好玩，我尝试让一堆 AI 模型来计算出这些数组的形状。这是结果： 我使用了这个查询：

执行以下 python 代码 A = np.ones((10,20,30,40)) i = np.array([1,2,3]) j = np.array([[0],[1]]) B = A[:,i,j,:] C = A[:,:,i,j] D = A[:,i,:,j] E = A[:,1:4,j,:] B、C、D 和 E 有什么形状？ Claude 3.7 使用了“扩展思维”。这是所有不正确的输出： AI | B | C | D | E
---|---|---|---|---
GPT 4.1 | 10×2×3×30
Grok 3 | 10×3×30×2 | 10×3×2×40
Claude 3 Opus | 10×3×2×30 | 10×20×3×2 | 10×3×30×2 | 10×3×2×40
Llama 4 Maverick | 10×3×30×2 | 10×3×2×40
o3 | 10×2×3×30
Claude 3.7 | 10×3×30×2 | 10×3×2×40
AI | B | C | D | E
---|---|---|---|---
GPT 4.1 | ✔️ | ✔️ | X | ✔️
Grok 3 | ✔️ | ✔️ | X | X
Claude 3 Opus | X | X | X | X
Llama 4 Maverick | ✔️ | ✔️ | X | X
o3 | ✔️ | ✔️ | X | ✔️
Claude 3.7 | ✔️ | ✔️ | X | X
Gemini 2.5 Pro | ✔️ | ✔️ | ✔️ | ✔️
DeepSeek R1 | ✔️ | ✔️ | ✔️ | ✔️
(DeepSeek 的思维链使用了“等待”76 次。它第一次就完全正确了，但是当我再次尝试时，它不知何故将 B、C 和 D 都弄错了，但是 E 却正确了。)

这太疯狂了。使用基本功能不应该需要解决疯狂的逻辑难题。

你可能会想，“好吧，我只会将自己限制在以简单的方式进行索引。”听起来不错，但有时你 需要 高级索引。即使你正在做一些简单的事情，你仍然需要小心避免疯狂的情况。

这再次使一切变得不可读。即使你只是在阅读以简单方式使用索引的代码，你怎么 知道 它是简单的？如果你看到 A[B,C]，那可能在做几乎任何事情。要理解它，你需要记住 A、B 和 C 的形状，并仔细考虑所有情况。而且，当然，A、B 和 C 通常是由 其他 代码产生的，你 也 需要考虑这些代码……

我不喜欢 NumPy 的函数

为什么 NumPy 最终会有一个如此令人困惑的 np.linalg.solve(A,B) 函数？我想象他们首先让它在 A 是一个 2D 数组，b 是一个 1D 或 2D 数组时工作，就像 A⁻¹b 或 A⁻¹B 的数学符号一样。

到目前为止，一切都很好。但是然后有人可能会拿出一个 3D 数组。如果你可以使用循环，解决方案将是“使用带有循环的旧函数”。但是你不能使用循环。所以基本上有三个选择：

  1. 他们可以添加一些额外的 axes 参数，以便用户可以指定要操作的维度。也许你可以编写 solve(A,B,axes=[[1,2],1])。
  2. 他们可以为不同的情况创建具有不同名称的不同函数。也许 solve_matrix_vector 会做一件事，solve_tensor_matrix 会做另一件事。
  3. 他们可以添加一个约定：对于 solve 将如何在内部尝试对齐维度做出一些任意选择。然后，用户有责任弄清楚并遵守这些约定。

所有这些选项都很糟糕，因为它们都无法真正处理存在组合数量的不同情况的事实。 NumPy 选择了：所有选项。有些函数具有 axes 参数。有些函数具有不同名称的不同版本。有些函数具有约定。有些函数具有约定 和 axes 参数。有些函数根本不提供任何向量化版本。

但是 NumPy 最大 的缺陷是：假设你创建了一个函数来解决一些具有给定形状的数组的问题。现在，你如何将它应用于一些更大数组的特定维度？答案是：你以一种更复杂的方式从头开始重写你的函数。编程的基本原则是抽象——解决简单的问题，然后将解决方案用作更复杂问题的构建块。 NumPy 不允许你这样做。

请注意

最后一个例子来向你展示我在说什么。每当我抱怨 NumPy 时，人们总是希望看到一个带有 self-attention 的例子，这是现代语言模型背后的核心技巧。所以好吧。这是一个实现，我谦虚地建议它比我搜索“self-attention numpy”时发现的所有 227 个版本都要好：

# self attention by your friend dynomight

input_dim = 4 
seq_len = 4 
d_k = 5 
d_v = input_dim 
 
X = np.random.randn(seq_len, input_dim) 
W_q = np.random.randn(input_dim, d_k) 
W_k = np.random.randn(input_dim, d_k) 
W_v = np.random.randn(input_dim, d_v) 
 
def softmax(x, axis): 
  e_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) 
  return e_x / np.sum(e_x, axis=axis, keepdims=True) 
 
def attention(X, W_q, W_k, W_v): 
  d_k = W_k.shape[1] 
  Q = X @ W_q 
  K = X @ W_k 
  V = X @ W_v 
  scores = Q @ K.T / np.sqrt(d_k) 
  attention_weights = softmax(scores, axis=-1) 
  return attention_weights @ V

这很好。一些 axis 的东西有点模糊，但没关系。

但是语言模型真正需要的是 multi-head attention，你可以在其中并行地进行几次 attention，然后合并结果。我们该怎么做呢？

首先，让我们想象一下，我们生活在一个允许我们使用抽象的理智世界中。然后，你可以简单地在一个循环中调用前面的函数：

# multi-head self attention by your friend dynomight
# if only we could use loops

n_head = 2
 
X = np.random.randn(seq_len, input_dim) 
W_q = np.random.randn(n_head, input_dim, d_k) 
W_k = np.random.randn(n_head, input_dim, d_k) 
W_v = np.random.randn(n_head, input_dim, d_v) 
W_o = np.random.randn(n_head, d_v, input_dim // n_head)
def multi_head_attention(X, W_q, W_k, W_v, W_o): 
  projected = [] 
  for n in range(n_head): 
    output = attention(X, W_q[n,:,:], W_k[n,:,:], W_v[n,:,:]) 
    my_proj = output @ W_o[n,:,:] 
    projected.append(my_proj) 
  projected = np.array(projected) 
 
  output = [] 
  for i in range(seq_len): 
    my_output = np.ravel(projected[:,i,:]) 
    output.append(my_output) 
  return np.array(output)

看起来很蠢，对吧？是的——谢谢你！聪明是不好的。

但是我们没有生活在一个理智的世界中。所以相反，你需要这样做：

# multi-head self attention by your friend dynomight
# all vectorized and bewildering

def multi_head_attention(X, W_q, W_k, W_v, W_o): 
  d_k = W_k.shape[-1] 
  Q = np.einsum('si,hij->hsj', X, W_q) 
  K = np.einsum('si,hik->hsk', X, W_k) 
  V = np.einsum('si,hiv->hsv', X, W_v) 
  scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_k) 
  weights = softmax(scores, axis=-1)
  output = weights @ V 
  projected = np.einsum('hsv,hvd->hsd', output, W_o) 
  return projected.transpose(1, 0, 2).reshape(seq_len, input_dim)

哈！哈哈哈哈哈哈！

那么该怎么办？

需要明确的是，我只是在说 NumPy 是“除了所有其他数组语言之外最糟糕的数组语言”。如果我没有更好的建议，抱怨又有什么意义呢？

嗯，实际上我确实有更好的建议。我制作了一个“更好”的 NumPy 的原型，我认为它保留了所有的力量，同时消除了所有的尖锐边缘。我以为这只是一个简短的动机性介绍，但是在我开始写作之后，邪恶控制了我，现在我们已经过了 3000 字了。

此外，在一个人的狂暴的论战和一个人建设性的数组语言 API 提案之间保持一些距离可能是明智的。所以我下次会介绍我的新东西。