Numpy Subtract two arrays of equal ndim but different shape

Question

So I have two ndarrays:

A with shape (N,a,a), a stack of N arrays of shape (a,a) basically

B with shape (8,M,a,a), a matrix of 8 x M arrays of shape (a,a)

I need to subtract B from A (AB) such that the resulting array is of shape (8,M*N,a,a). More verbosely each (M total) of the 8 arrays of B needs to be subtracted from each array in A, resulting in 8*M*N subtractions between (a,a) shape arrays.

How can I do this in a vectorized manner without loops? This thread does something similar but in lower dimensions and I can't figure out how to extend it.

Answer 1

A = np.arange(8).reshape(2,2,2)
B = np.ones(shape=(8,4,2,2))

General broadcasting works if dimensions are the same or if one dimension is 1, so we do this;

a = A[np.newaxis, :, np.newaxis, :, :]
b = B[:, np.newaxis, :, :, :]

a.shape  # <- (1,2,1,2,2)
b.shape  # <- (8,1,4,2,2)

Now when you can do broadcasting

c = a - b
c.shape  # <- (8,2,4,2,2)

And when you reshape the (2x4=8) components get aligned.

c.reshape(8,-1,2,2) 

The ordering of the new axes dictates the reshaping, so be careful with that.