【学习笔记】einsum（爱因斯坦求和）

文章目录

• einsum（爱因斯坦求和）
• • 基本语法
  • 对角线操作
  • 矩阵乘法
  • 向量点积
  • 张量积
• 总结

---

einsum（爱因斯坦求和）

einsum 是 Python 中 NumPy 库的一个函数，它提供了一种非常灵活的方式来进行多维数组的对角线操作、矩阵乘法、向量点积、张量积等操作。不过这需要一定的学习成本。本文旨在了解einsum的使用方法，并给出一些演示案列。

基本语法

numpy.einsum(subscripts, *operands)

• subscripts 是一个字符串，用来描述操作的输入和输出数组的维度以及如何进行操作。
• *operands 是一个或多个数组，即要进行操作的输入数组。

对角线操作

# 对角线操作：
import numpy as np
a = np.arange(9).reshape(3, 3)
print(a)
print("-"*20)
np.einsum('ii', a)  

打印结果为：

矩阵乘法

# 矩阵乘法
a = np.arange(4).reshape(2, 2)
b = np.arange(4, 8).reshape(2, 2)
print(a)
print("="*20)
print(b)
np.einsum('ij,jk->ik', a, b)

打印结果为：

向量点积

# 向量点积：
a = np.array([1, 2, 3])
b = np.array([0, 1, 0])
print(a)
print("="*20)
print(b)
np.einsum('i,i', a, b)  # 结果是 2

打印结果为：

张量积

# 张量积：
a = np.array([1, 2])
b = np.array([3, 4])
print(a)
print("="*20)
print(b)
np.einsum('i,j', a, b)  # 结果是 [[ 3.,  4.], [ 6.,  8.]]

总结

einsum 非常强大，但也需要仔细设计输入的 subscripts 字符串，以确保正确执行所需的操作。