torch.einsum(equation_str, *tensors)

import torch
import numpy as np

x = torch.rand((2,3))

# permutation of tensors
torch.einsum("ij->ji", x).size()

res: torch.Size([3, 2])

# summation
torch.allclose(torch.einsum("ij->", x), x.sum())

res: True

# summing all rows, which returns an array with the shape of columns
torch.allclose(torch.einsum('ij->j',x), x.sum(0))

res: True

# summing all columns, which returns an array with the shape of rows
torch.allclose(torch.einsum('ij->i',x), x.sum(1))

res: True

# matrix vector multiplication
v = torch.rand((1,3))

print(torch.einsum("ij,kj->ik", x,v))
print('--'*20)
print(x@v.T)

res:

tensor([[0.5172],
    [0.3321]])
----------------------------------------
tensor([[0.5172],
    [0.3321]])

# matrix matrix multiplication

print(torch.einsum("ij,kj->ik", x,x))
print("--"*20)
print(x@x.T)

res:

tensor([[0.6382, 0.4915],
            [0.4915, 0.9763]])
----------------------------------------
tensor([[0.6382, 0.4915],
        [0.4915, 0.9763]])

# dot product of all the rows of the matrix with itself

print(torch.einsum("ij,ij->i", x,x))

res: tensor([0.6382, 0.9763])

# even the elementwise multiplication

print(torch.einsum("ij,ij->ij", x,x))
print("--"*20)
print(torch.multiply(x,x))

res:
tensor([[0.4260, 0.0994, 0.1127],
        [0.0325, 0.8927, 0.0511]])
----------------------------------------
tensor([[0.4260, 0.0994, 0.1127],
        [0.0325, 0.8927, 0.0511]])

# outer product
x = torch.rand((3))
y = torch.rand((5))

torch.einsum("i,j -> ij", a,b)

res:
tensor([[0.0897, 0.2208, 0.3084, 0.0471, 0.0190],
            [0.0379, 0.0933, 0.1304, 0.0199, 0.0080],
            [0.0608, 0.1495, 0.2088, 0.0319, 0.0128]])

# batch matrix multiplication 

a = torch.rand((3,2,5))
b = torch.rand((3,5,3))

# you can do this with bmm
torch.einsum("ijk, ikl-> ijl", a,b)

res:
tensor([[[1.0678, 1.2429, 1.1200],
             [0.3129, 0.3637, 0.7859]],
    
            [[1.3433, 0.4159, 1.4224],
             [1.9831, 0.8470, 1.6947]],
    
            [[1.4838, 2.1888, 1.6716],
             [0.9673, 1.3760, 0.7142]]])

# matrix diagonal
x = torch.rand((3,3))

torch.einsum("ii->i", x)

res:
tensor([0.9835, 0.6807, 0.7448])

x

res:
tensor([[0.9835, 0.9617, 0.7266],
            [0.2494, 0.6807, 0.9855],
            [0.8947, 0.6352, 0.7448]])