tf.einsum

Свертка тензора по указанным индексам и внешнее произведение.

tf.einsum(
    equation, *inputs, **kwargs
)

Einsum позволяет определять тензоры, определяя их элементное вычисление. Это вычисление определяется equation, краткой формой, основанной на суммировании по Эйнштейну. В качестве примера рассмотрим умножение двух матриц A и B для формирования матрицы C. Элементы C задаются:

\[ C_{i,k} = \sum_j A_{i,j} B_{j,k} \]

или

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

Соответствующая строка einsum equation:

ij,jk->ik

В общем случае, чтобы преобразовать элементное уравнение в строку equation , используйте следующую процедуру (промежуточные строки для примера умножения матриц приведены в скобках):

1. удалить имена переменных, скобки и запятые, (ik = sum_j ij * jk)
2. заменить "*" на ",", (ik = sum_j ij , jk)
3. удалить знаки суммирования и (ik = ij, jk)
4. переместить выходной результат вправо, заменив "=" на "->". (ij,jk->ik)

Примечание: Если индексы вывода не указаны, повторяющиеся индексы суммируются. Таким образом, ij,jk->ik можно упростить до ij,jk.

Многие распространенные операции могут быть выражены таким способом. Например:

Умножение матриц

m0 = tf.random.normal(shape=[2, 3])
m1 = tf.random.normal(shape=[3, 5])
e = tf.einsum('ij,jk->ik', m0, m1)
# output[i,k] = sum_j m0[i,j] * m1[j, k]
print(e.shape)
(2, 5)

Повторяющиеся индексы суммируются, если индексы вывода не указаны.

e = tf.einsum('ij,jk', m0, m1)  # output[i,k] = sum_j m0[i,j] * m1[j, k]
print(e.shape)
(2, 5)

Скалярное произведение

u = tf.random.normal(shape=[5])
v = tf.random.normal(shape=[5])
e = tf.einsum('i,i->', u, v)  # output = sum_i u[i]*v[i]
print(e.shape)
()

Внешнее произведение

u = tf.random.normal(shape=[3])
v = tf.random.normal(shape=[5])
e = tf.einsum('i,j->ij', u, v)  # output[i,j] = u[i]*v[j]
print(e.shape)
(3, 5)

Транспонирование

m = tf.ones(2,3)
e = tf.einsum('ij->ji', m0)  # output[j,i] = m0[i,j]
print(e.shape)
(3, 2)

Diag

m = tf.reshape(tf.range(9), [3,3])
diag = tf.einsum('ii->i', m)
print(diag.shape)
(3,)

След

# Repeated indices are summed.
trace = tf.einsum('ii', m)  # output[j,i] = trace(m) = sum_i m[i, i]
assert trace == sum(diag)
print(trace.shape)
()

Умножение матриц с батчем

s = tf.random.normal(shape=[7,5,3])
t = tf.random.normal(shape=[7,3,2])
e = tf.einsum('bij,bjk->bik', s, t)
# output[a,i,k] = sum_j s[a,i,j] * t[a, j, k]
print(e.shape)
(7, 5, 2)

Этот метод не поддерживает векторизацию по именованным осям. Все оси с совпадающими метками должны иметь одинаковую длину. Если у вас есть оси длиной 1, используйте tf.squeeze или tf.reshape, чтобы их устранить.

Чтобы написать код, не зависящий от количества индексов во входных данных, используйте многоточие. Многоточие является заполнительм для "любых других индексов, которые подходят сюда".

Например, чтобы выполнить умножение матриц в стиле NumPy с векторизацией для батчей, где умножение матриц действует на последних двух осях входных данных, используйте:

s = tf.random.normal(shape=[11, 7, 5, 3])
t = tf.random.normal(shape=[11, 7, 3, 2])
e =  tf.einsum('...ij,...jk->...ik', s, t)
print(e.shape)
(11, 7, 5, 2)

Einsum будет векторизоваться по осям, охваченным многоточием.

s = tf.random.normal(shape=[11, 1, 5, 3])
t = tf.random.normal(shape=[1, 7, 3, 2])
e =  tf.einsum('...ij,...jk->...ik', s, t)
print(e.shape)
(11, 7, 5, 2)