Add masked matrix tests, ensure mask.shape remains correct