einsum(爱因斯坦求和)
在Python中,特别是使用NumPy或PyTorch等科学计算库时,爱因斯坦求和约定(Einstein summation convention)是一种强大的工具,它允许你以简洁的方式指定多维数组(张量)之间的操作,如求和、乘法和广播。虽然Python本身并不直接支持爱因斯坦求和约定的语法,但一些库(如NumPy的einsum
函数和PyTorch的einsum
操作)提供了对它的支持。
NumPy中的einsum
NumPy的einsum
函数允许你指定数组操作的计算方式,通过一种紧凑的字符串表示法。这个字符串描述了输入数组的形状、操作的维度以及如何对它们进行求和或广播。
基本语法
numpy.einsum(subscripts, *operands, **kwargs)
subscripts
:一个字符串,描述了操作。它指定了输入数组的每个维度以及如何在输出数组中组合它们。operands
:一个或多个数组,是操作的输入。**kwargs
:其他关键字参数,如out
(指定输出数组),dtype
(指定输出数组的数据类型),order
(指定数组的内存布局)等。
示例
- 矩阵乘法
import numpy as np
a = np.random.rand(3, 4)
b = np.random.rand(4, 5)
# 使用einsum进行矩阵乘法
c = np.einsum('ij,jk->ik', a, b)
# 等同于使用numpy的dot函数
# c = np.dot(a, b)
在这个例子中,'ij,jk->ik'
指定了操作:i
和j
是a
的维度,j
和k
是b
的维度,输出c
的维度是i
和k
,其中j
维度被求和(这是爱因斯坦求和约定的核心)。
- 计算内积
a = np.random.rand(3)
b = np.random.rand(3)
# 使用einsum计算内积
dot_product = np.einsum('i,i->', a, b)
# 等同于使用numpy的dot函数(对于一维数组)
# dot_product = np.dot(a, b)
在这个例子中,'i,i->'
指定了操作:i
是a
和b
的公共维度,输出是一个标量(没有维度),因为所有维度都被求和了。
PyTorch中的einsum
PyTorch也提供了一个einsum
操作,其语法和功能与NumPy的einsum
非常相似。
示例
import torch
a = torch.randn(3, 4)
b = torch.randn(4, 5)
# 使用einsum进行矩阵乘法
c = torch.einsum('ij,jk->ik', [a, b])
# 等同于使用torch的matmul函数
# c = torch.matmul(a, b)
在这个PyTorch的例子中,einsum
的调用方式与NumPy非常相似,但是注意输入数组是作为列表传递给einsum
的。
总结
爱因斯坦求和约定是一种强大的工具,它允许你以一种非常紧凑和灵活的方式指定多维数组之间的操作。虽然Python本身不直接支持这种语法,但NumPy和PyTorch等库提供了对它的支持,使得你可以利用这种约定来编写更简洁、更高效的代码。