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)

Диагональ

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)