[MaskedArray] mask.shape diverges from data.shape

After the following calculation with MaskedArray, the shape of the mask diverges from the shape of the data. While the shape of data is correct.

Reproducing code example:

import numpy as np
import numpy.ma as ma

n_particles = 5
max_n_features = 4
feature_mask = np.random.randint(0,2,size=(max_n_features, n_particles)).astype(bool)

mask_state = feature_mask.reshape((max_n_features, n_particles, 1, 1))
mask_state = np.repeat(mask_state, repeats=3, axis=2) # (4,5,3,1)
mask_covariance = np.repeat(mask_state, repeats=3, axis=3) # (4,5,3,3)

Q = ma.MaskedArray(np.random.uniform(0,1,size=(max_n_features, n_particles,3,3)), mask = mask_covariance)

surprise = ma.MaskedArray(np.random.uniform(0,1,size=(max_n_features, n_particles,3,1)), mask = mask_state)
exp_term = np.exp(-0.5 * np.transpose(surprise, axes=(0, 1, 3, 2)) @ np.linalg.inv(Q) @ surprise)

print(exp_term.data.shape, exp_term.mask.shape) # prints: (4, 5, 1, 1) (4, 5, 3, 3)
print(exp_term.data.shape == exp_term.mask.shape) # prints: False which seems to be wrong.

Error message:

No exception or error message is raised. but the last line print(exp_term.data.shape == exp_term.mask.shape)
prints False. In this calculation the shape of exp_term.data is the correct one. and exp_term.mask has wrong shape.

When the shapes are different, it does not allow to reshape or squeeze or even print the MaskedArray.

exp_term.data.shape is (4,5,1,1) while exp_term.mask.shape is (4,5,3,3). but nice thing about mask is that for every i and j the exp_term.mask[i, j, :,:] are either all True or all False. Or in other words, the following is true

np.all(   np.any(exp_term.mask,axis=(2,3)) == np.all(exp_term.mask,axis=(2,3))   )

So there is a workaround:

new_mask = np.expand_dims(np.all(exp_term.mask,axis=(2,3)),axis=(2,3))
exp_term = ma.MaskedArray(exp_term.data, mask = new_mask) 

NumPy/Python version information:

numpy version: 1.20.3 

python version: 3.8.5 (default, Feb 16 2021, 00:08:37) 
[GCC 9.3.0]

Start by running the provided MaskedArray reproducer with the reported NumPy version and inspect the transpose, matrix multiplication, inverse, and mask propagation involved in the expression. Trace the MaskedArray operation that leaves data and mask with different shapes, then add a regression test for the reproducer and confirm that the resulting shapes remain consistent.