import torch
import numpy as np
a = torch.rand(32, 48, 9, 3)
print(np.einsum('...ji,...jk ->...ik', a, a).shape)
# (32, 48, 3, 3)
print(torch.einsum('...ji,...jk ->...ik', [a, a]).shape)
# RuntimeError: ellipsis must represent 2 dimensions in all terms
Found during porting matting laplacian to PyTorch: https://github.com/NVIDIA/FastPhotoStyle/blob/master/photo_smooth.py#L76