You can use:
mse = ((A - B)**2).mean(axis=ax)
Or
mse = (np.square(A - B)).mean(axis=ax)
- with
ax=0
the average is performed along the row, for each column, returning an array - with
ax=1
the average is performed along the column, for each row, returning an array - with
ax=None
the average is performed element-wise along the array, returning a scalar value