Grokking Einstein summation with NumPy.

grok_einsum.py

import numpy as np

A = np.array([[1, 1, 1], [2, 2, 2], [5, 5, 5]])
B = np.array([[0, 1, 0], [1, 1, 0], [1, 1, 1]])

(i_s, j_s, k_s) = (len(A), len(A[0]), len(B[0]))
for do_sum in [True, False]:
    if do_sum:
        C = np.zeros((i_s, k_s))
        print(np.einsum("ij,jk->ik", A, B))
    else:
        C = np.zeros((i_s, j_s, k_s))
        print(np.einsum("ij,jk->ijk", A, B))

    for i in range(i_s):
        for j in range(j_s):
            for k in range(k_s):
                C_ijk = A[i, j] * B[j, k]
                if do_sum:
                    C[i, k] += C_ijk
                else:
                    C[i, j, k] = C_ijk

    print(C)