Notice that the shapes you are using for “mask” matrices are derived from the bias vectors, so you will be treating each column of A the same way because of “broadcasting”. That’s not what the instructions intended: they tell you to make the shape of the mask directly the shape of the corresponding A value, which means that each sample will be treated differently in terms of which neurons are dropped.
This is an interesting point to consider: does it make sense to treat each sample the same within a given minibatch or not? This topic has come up a number of times before and here’s a thread which discusses this in more detail and actually shows some experiments comparing the results of the two methods of implementing dropout.
Regardless of the conclusion of any research on the alternatives, in order to pass the tests here you have to make your code conform to the “each sample is handled differently” strategy.