Умножьте тензоры, содержащие матрицы, следуя правилу умножения матриц

Вероятно, не самое эффективное, но вот ручное решение:

M = np.random.randint(0, 10, (2, 2, 2, 2))

def matmatmul(a, b):
    
    output = np.zeros((a.shape[0], b.shape[1]), dtype = object)
    
    for i in range(output.shape[0]):
        for j in range(output.shape[1]):
            row = a[i]
            col = b[:, j]
            
            output[i,j] = sum([r @ c for r,c in zip(row, col)])
            
    return output

def matmatpow(a, n):
    if n == 1:
        return a
    else:
        output = matmatmul(a, a)
        
        for i in range(2, n):
            output = matmatmul(output, a)
        
        return output
    
M2 = matmatpow(M, 2)

print(M2)

[[A, B], [C, D]] = M

assert np.all(M2[0,0] == A@A + B@C)
assert np.all(M2[0,1] == A@B + B@D)
assert np.all(M2[1,0] == C@A + D@C)
assert np.all(M2[1,1] == C@B + D@D)

Определение набора массивов (2,2) и их составных частей:

In [45]: A,B,C,D = [np.arange(i,i+4).reshape(2,2) for i in range(4)]    
In [46]: M=np.array([[A,B],[C,D]])

Ваш желаемый массив M^2:

In [47]: np.array([[A@A + B@C, A@B + B@D],
    ...:       [C@A + D@C, C@B + D@D]])
Out[47]: 
array([[[[12, 16],
         [28, 40]],

        [[16, 20],
         [40, 52]]],


       [[[28, 40],
         [44, 64]],

        [[40, 52],
         [64, 84]]]])

То же самое с использованием einsum. В этом j и l представлены измерения суммы продуктов:

In [48]: np.einsum('ijkl,jmln->imkn',M,M)
Out[48]: 
array([[[[12, 16],
         [28, 40]],

        [[16, 20],
         [40, 52]]],


       [[[28, 40],
         [44, 64]],

        [[40, 52],
         [64, 84]]]])

matmul эквивалентно ijkl,ijlm->ijkm, где ij — batch размеры, а l — сумма произведений. Часто einsum можно воспроизвести с некоторым изменением формы и общим транспонированием. Но я оставлю это для изучения другим.

Играя с индексами einsum, транспонируя и изменяя массивы, я могу получить эквивалент:

In [56]: np.matmul(M.transpose(0,2,1,3).reshape(4,4),M.transpose(0,2,1,3).reshape(4,4))
Out[56]: 
array([[12, 16, 16, 20],
       [28, 40, 40, 52],
       [28, 40, 40, 52],
       [44, 64, 64, 84]])

который при еще немного массировании становится желаемым (4,4,4,4)

In [57]: np.matmul(M.transpose(0,2,1,3).reshape(4,4),M.transpose(0,2,1,3).reshape(4,4)).reshape(2,2,2,2).transpose(0,2,1,3)
Out[57]: 
array([[[[12, 16],
         [28, 40]],

        [[16, 20],
         [40, 52]]],


       [[[28, 40],
         [44, 64]],

        [[40, 52],
         [64, 84]]]])

Вы просто вычисляете нормальное матричное произведение блочной матрицы 4x4, созданной путем соединения меньших матриц от A до D.

В MATLAB ваш ожидаемый результат с использованием некоторых произвольных матриц:

A = [1, 2
     3, 4];
B = [5, 6
     7, 8];
C = [ 9, 10
     11, 12];
D = [13, 14
     15, 16];

res = [A*A + B*C, A*B + B*D
       C*A + D*C, C*B + D*D]

res =
   118   132   174   188
   166   188   254   276
   310   356   494   540
   358   412   574   628

Блочная матрица 4x4 и ее квадрат:

M = [A, B
     C, D];
res2 = M^2

res2 =
   118   132   174   188
   166   188   254   276
   310   356   494   540
   358   412   574   628