the same random index at each evaluation
But it seems like ridx gets evaluated multiple times, namely when indexing each of the arrays.
No, I would only minibatch the mask as that already controls the indexing to the data and mu.
Then I totally do not understand how that should work. maskshared has the same shape as mask, so n*8*2. Now, maskshared[ridx, :, :] has shape 500*8*2, right? But data still has n*8*2, so how can we index data with maskshared[ridx, :, :]?