einsum的使用

### 使用 `einsum` 的基本概念 `einsum` 是一种强大的函数，可以用来执行各种张量操作，比如矩阵乘法、点积、外积以及更复杂的多维数组运算。它通过爱因斯坦求和约定来描述这些操作。 #### NumPy 中的 `einsum` NumPy 提供了一个灵活的 `einsum` 函数，允许用户定义任意维度上的张量计算。其语法简单明了，能够处理多种线性代数运算[^1]。 ```python import numpy as np # 创建两个二维数组 a = np.array([[1, 2], [3, 4]]) b = np.array([[5, 6], [7, 8]]) # 矩阵乘法 result = np.einsum('ij,jk->ik', a, b) print(result) ``` 上述代码实现了标准的矩阵乘法操作，其中 `'ij,jk->ik'` 表示第一个输入张量沿第 i 和 j 轴展开，第二个输入张量沿第 j 和 k 轴展开，最终结果是一个新的张量沿着轴 i 和 k 组合而成。 --- #### PyTorch 中的 `einsum` PyTorch 同样支持类似的 `einsum` 功能，适用于 GPU 或 CPU 上的操作。它的实现方式与 NumPy 类似，但在性能上可能有所不同。 ```python import torch # 创建两个二维张量 a = torch.tensor([[1., 2.], [3., 4.]]) b = torch.tensor([[5., 6.], [7., 8.]]) # 执行矩阵乘法 result = torch.einsum('ij,jk->ik', a, b) print(result) ``` 此代码片段展示了如何利用 PyTorch 的 `einsum` 进行相同的矩阵乘法操作。值得注意的是，在某些情况下，TensorFlow 可能提供额外的功能而 PyTorch 尚未完全覆盖。 --- #### TensorFlow 中的 `einsum` TensorFlow 的 `einsum` 实现也遵循同样的逻辑框架，并且提供了更多的灵活性和支持功能[^2]。 ```python import tensorflow as tf # 创建两个二维张量 a = tf.constant([[1., 2.], [3., 4.]]) b = tf.constant([[5., 6.], [7., 8.]]) # 执行矩阵乘法 result = tf.einsum('ij,jk->ik', a, b) print(result.numpy()) ``` 尽管两者在核心功能上有重叠之处，但 TensorFlow 在特定领域（如快速傅里叶变换等）仍具有独特优势。 --- ### 总结对比 | 特性 | NumPy | PyTorch | TensorFlow | |---------------------|-----------------------|----------------------|----------------------| | 支持设备 | CPU | CPU/GPU | CPU/GPU/TPU | | 额外特性 | 基础数学运算 | 更强动态图能力 | 更丰富的工具集 | 以上表格总结了三种库对于 `einsum` 的主要差异及其适用场景。 ---