enter code hereI have a ndarray structure with 4 dimensions: (2x2xNxK). You can imagine this as K different stacks, containing each of them N matrices of dimension 2×2.
For each stack, I’m trying to compute the matrix product of its N matrices (without For loop). So in the end, I should have K matrices of dimensions 2×2 (the product of N 2×2 matrices remains a 2×2 matrix).

If my array was of dimensions (2x2xN), this is quite straightforward to perform matrix product along the 3rd axis with the following function:
A_total = np.linalg.multi_dot(A) #— A being a (2x2xN) array

But with a (2x2xNxK) structure, I don’t manage to perform the same operation along the 3rd axis, while ignoring the 4th axis.

Would you have any suggestion on this? I have tried with np.einsum() and np.tensordot() but don’t quite understand how to correctly use those functions.