TensorFlow如何实现张量收缩

tf.einsum

einsum(
    equation,
    *inputs
)

定义在：tensorflow/python/ops/special_math_ops.py.

参考指南：数学>减少张量的计算

任意维度张量之间的广义收缩.

这个函数返回一个张量,其元素其元素是由等式定义的,这是由爱因斯坦求和公式所启发的速写形式定义的.作为示例,考虑将两个矩阵 A 和 B 相乘以形成矩阵C.C的元素由下式给出：

C[i,k] = sum_j A[i,j] * B[j,k]

相应的等式是：

ij,jk->ik

一般来说, 方程是从较熟悉的元素方程得到：

1. 删除变量名称、括号和逗号；
   
2. 用 "*" 替换 ","；
   
3. 删除总和标志；
   
4. 将输出移到右侧,并将 "=" 替换为 "->>".
   

许多常见操作可以用这种方式来表示.例如:

# Matrix multiplication
>>> einsum('ij,jk->ik', m0, m1)  # output[i,k] = sum_j m0[i,j] * m1[j, k]

# Dot product
>>> einsum('i,i->', u, v)  # output = sum_i u[i]*v[i]

# Outer product
>>> einsum('i,j->ij', u, v)  # output[i,j] = u[i]*v[j]

# Transpose
>>> einsum('ij->ji', m)  # output[j,i] = m[i,j]

# Batch matrix multiplication
>>> einsum('aij,ajk->aik', s, t)  # out[a,i,k] = sum_j s[a,i,j] * t[a, j, k]

此函数的行为类似于 numpy.einsum,但不支持：

• 椭圆(下标像：ij...,jk...->ik...)
• 一个轴在单个输入上出现多次的下标(例如 ijj、k->> ik).
• 在多个输入之间求和的下标(例如 ij、ij、jk->> ik).

ARGS：

• equation：一种描述收缩的 str,与 numpy.einsum 格式相同.
• * inputs：合同的输入(每个张量),其形状应与 equation 一致.

返回：

返回收缩的张量,形状由 equation 决定.

注意：

• ValueError：如果 equation 格式不正确,equation 隐含的输入数与 len(inputs) 不匹配,一个轴出现在输出下标中,但不显示在任何输入中,输入的维数与其下标中的索引数不同,或者输入形状沿特定轴线不一致.