tf.einsum

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

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

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

или

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.squeseze или 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)