TensorFlow如何實(shí)現(xiàn)張量收縮

tf.einsum

einsum(
    equation,
    *inputs
)

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

參考指南：數(shù)學(xué)>減少?gòu)埩康挠?jì)算

任意維度張量之間的廣義收縮.

這個(gè)函數(shù)返回一個(gè)張量,其元素其元素是由等式定義的,這是由愛(ài)因斯坦求和公式所啟發(fā)的速寫形式定義的.作為示例,考慮將兩個(gè)矩陣 A 和 B 相乘以形成矩陣C.C的元素由下式給出：

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

相應(yīng)的等式是：

ij,jk->ik

一般來(lái)說(shuō), 方程是從較熟悉的元素方程得到：

1. 刪除變量名稱、括號(hào)和逗號(hào)；
   
2. 用 "*" 替換 ","；
   
3. 刪除總和標(biāo)志；
   
4. 將輸出移到右側(cè),并將 "=" 替換為 "->>".
   

許多常見(jiàn)操作可以用這種方式來(lái)表示.例如:

# 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]

此函數(shù)的行為類似于 numpy.einsum,但不支持：

• 橢圓(下標(biāo)像：ij...,jk...->ik...)
• 一個(gè)軸在單個(gè)輸入上出現(xiàn)多次的下標(biāo)(例如 ijj、k->> ik).
• 在多個(gè)輸入之間求和的下標(biāo)(例如 ij、ij、jk->> ik).

ARGS：

• equation：一種描述收縮的 str,與 numpy.einsum 格式相同.
• * inputs：合同的輸入(每個(gè)張量),其形狀應(yīng)與 equation 一致.

返回：

返回收縮的張量,形狀由 equation 決定.

注意：

• ValueError：如果 equation 格式不正確,equation 隱含的輸入數(shù)與 len(inputs) 不匹配,一個(gè)軸出現(xiàn)在輸出下標(biāo)中,但不顯示在任何輸入中,輸入的維數(shù)與其下標(biāo)中的索引數(shù)不同,或者輸入形狀沿特定軸線不一致.