Использование einsum для транспонирования времен матрицы транспонирования: x@A@x^T

Это эквивалент этого цикла:

result_einsum = torch.einsum('ki,ij,kj->k', B, A, B)[..., None]

Часть [..., None] существует для добавления дополнительного измерения к выводу. Не удалось понять, как это сделать с чистым einsum, хотя дополнительная часть, по крайней мере, не имеет копирования.

С (m,n) и (n,n)

res = np.einsum('ij,jk,lk-il', B,A,B)

должен создать массив (m,m).

'ij,jk,ik->i' должно составить (m,)

С matmul/@ я думаю, это даст

B@A@(B.T) # (m,m)

B[:,None,:]@A@B[:,:,None] # (m,1,1)

Пишу с телефона, поэтому проверить сейчас не могу.

Использование einsum и tensordot:

import torch
import numpy as np

A = torch.tensor([[1.0, 0.0, 0.0, 0.0],
                  [0.0, 2.0, 0.0, 0.0],
                  [0.0, 0.0, 3.0, 0.0],
                  [0.0, 0.0, 0.0, 4.0]])

B = torch.tensor([[1.0, 4.0, 7.0, 10.0],
                  [2.0, 5.0, 8.0, 11.0],
                  [3.0, 6.0, 9.0, 12.0]])

# Using einsum

res_using_einsum = torch.einsum('bi,ij,bj -> b',B,A,B).unsqueeze(dim = -1)
print(res_using_einsum)
'''
tensor([[580.],
        [730.],
        [900.]])
'''

# Using tensordot

BA = torch.tensordot(B, A, dims = ([1],[0]))
'''
BA :
tensor([[ 1.,  8., 21., 40.],
        [ 2., 10., 24., 44.],
        [ 3., 12., 27., 48.]])
'''
res_using_tensordot = torch.tensordot(BA, B, dims =([1],[1]))#.unsqueeze(-1)

'''
res_using_tensordot :

tensor([[580., 650., 720.],
        [650., 730., 810.],
        [720., 810., 900.]])
'''
diagonal_result = torch.diagonal(res_using_tensordot, 0).unsqueeze(1)
'''
tensor([[580.],
        [730.],
        [900.]])

'''