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)