Expanded Knowledge

Dot Product (Vector Result) \( C_i = \sum_j A_{ij} B_j \)

import numpy as np
a = np.arange(0, 12).reshape(3, 4)
b = np.arange(0, 4)
c = np.einsum('ij,j->i', a, b)

Full Scalar Contraction \( C = \sum_i \sum_j A_{ij} B_{ij} \)

import numpy as np
a = np.arange(0, 12).reshape(3, 4)
b = np.arange(0, 12).reshape(3, 4)
c = np.einsum('ij,ij->', a, b)

Matrix Multiplication \( C_{ij} = \sum_k A_{ik} B_{kj} \)

import numpy as np
a = np.arange(0, 12).reshape(3, 4)
b = np.arange(0, 12).reshape(4, 3)
c = np.einsum('ik,kj->ij', a, b)

High-Order Tensor Contraction \( C_{ijkmno} = \sum_l A_{ijklno} B_{ijlmno} \)

import torch
A = torch.randn(16, 8, 5, 128, 128)
B = torch.randn(16, 8, 5, 128, 128)
A = A.unsqueeze(3)
B = B.unsqueeze(2)
C = torch.einsum('ijklno,ijlmno->ijkmno', A, B)