[Python] Efficiently compute all multi-dimensional traces for all offsets and store in matrix

I have a $N\times N \times N$ array $a$ and would like to implement the formula $$ b_{ij} = \sum_k a_{i+k,j+k,k} $$ efficiently.

Right now, I'm doing this via

b = np.zeros((N, N))
for i in range(N):
for j in range(N):
b[i,j] = np.einsum('iii', a[i:,j:,])


which seems quite inefficient.

Is there a way to do this without cython which works just with numpy (or any other, numpy-compatible interface such as jax.numpy)?

Continue reading...