Numpy中能否实现先执行惰性broadcasting再做reduce操作？

原生NumPy没有内置的惰性广播机制，你遇到的广播后归约产生大体积中间数组的问题，目前可以通过以下几种方案解决，兼顾内存效率和代码简洁性：

• 通过数学公式改写向量化逻辑，避免生成冗余中间数组
  针对两两MSE计算的场景，可以将MSE公式展开为平方和加和减两倍内积的形式，完全规避生成(N,N,d)维度的中间数组，内存占用仅为(N,N)级别，运行速度也远高于手动Python循环：

  def pairwise_mse(vectors):
      square_norm = np.sum(vectors ** 2, axis=1)
      # 利用广播和矩阵乘法实现，无d维度中间数组生成
      return square_norm[:, None] + square_norm[None, :] - 2 * vectors @ vectors.T
  

• 使用Numba JIT编译优化手动循环
  如果你不想修改原有mse函数的实现逻辑，可以用Numba装饰器编译相关函数，编译后的Python循环会被优化为机器码执行，速度和原生NumPy操作相当，同时也不会生成冗余的大中间数组：

  import numba
  import numpy as np
  
  @numba.njit
  def mse(x, y):
      return np.mean((x - y)**2, axis=-1)
  
  @numba.njit
  def pairwise_mse_numba(vectors):
      N = vectors.shape[0]
      res = np.empty((N, N), dtype=vectors.dtype)
      for i in range(N):
          res[i] = mse(vectors[i:i+1], vectors)
      return res
  

• 用Dask数组实现惰性计算
  Dask数组完全兼容NumPy API，默认采用惰性计算模式，会自动优化计算图避免生成不必要的中间数组，完全保留你习惯的广播语法，不需要修改原有函数实现：

  import dask.array as da
  import numpy as np
  
  N, d = 100, 3
  vectors = np.random.uniform(size=(N, d))
  # 转换为Dask数组，可自定义分块大小控制内存占用
  dask_vectors = da.from_array(vectors, chunks=(10, d))
  
  # 写法和原生NumPy完全一致，此时不会实际执行计算
  pairwise = mse(dask_vectors[None, :, :], dask_vectors[:, None, :])
  # 触发实际计算得到结果
  pairwise = pairwise.compute()
  

内容的提问来源于stack exchange，提问作者a_guest