Randomly masking a variable number of selected positions in a 2D tensor

username: Norix
cooked: <p>Hi.</p>
<p>I have a 2D tensor and a corresponding 2D boolean mask. The mask is telling me which positions are valid for being randomly masked. All other values in the tensor need to stay the same.</p>
<p>In the next step I want to randomly select 1 to N elements from each row of my tensor and change their values to a “mask value”, with N being the number of valid elements in that row.</p>
<p>I have some code to do this for each row individually, but I wonder if there is an efficient way to do this for all rows at once?</p>
<p>Can the following code be adapted to work with 2D tensors instead of 1D tensors?<br>
(“tensor” is my 1D tensor, “mask” is the corresponding 1D boolean mask of valid positions)</p>
<pre><code class="lang-auto"># get indices of True values in mask
mask_idx = mask.nonzero().flatten()

N = mask_idx.numel()  # number of valid positions
M = random.randint(1, N)  # number of positions to mask

# randomly select M out of N valid positions
idx_perm = torch.randperm(N)[:M]

# change mask to the randomly selected valid positions
mask.fill_(False)
mask[mask_idx[idx_perm]] = True

# mask values in tensor
tensor[mask] = mask_value
</code></pre>
<p>In addition, I do need to know the values for N and M for each row. Though those can also be calculated based on the row-sums of the corresponding mask before and after running the previous code.</p>
<p>Thank you.</p>

username: KFrank
cooked: <p>Hi Norix!</p>
<aside class="quote no-group quote-modified" data-username="Norix" data-post="1" data-topic="215554" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/n/9f8e36/48.png" class="avatar"> Norix:</div>
<blockquote>
<p>I have a 2D tensor and a corresponding 2D boolean mask. The mask is telling me which positions are valid for being randomly masked.<br>
…<br>
I want to randomly select 1 to N elements from each row of my tensor and change their values to a “mask value”, with N being the number of valid elements in that row.<br>
…<br>
I wonder if there is an efficient way to do this for all rows at once?</p>
</blockquote>
</aside>
<p>Yes, with a proviso*.</p>
<p>Here is a loop-free script that works on the entire two-dimensional<br>
tensors all at once:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2025)

shape = (5, 8)

t = torch.zeros (5, 8)
mask_value = 999

mask = torch.randint (2, (5, 8)).bool()                                                     # mask for valid mask location
print ('mask = ...')
print (mask)

ma = torch.randint (2, (5, 8)).bool()                                                       # locations to mask
mb = torch.nn.functional.one_hot ((torch.rand (5, 8) * mask).argmax (1), shape[1]).bool()   # mask at least one valid location

mask2 = (ma + mb) * mask                                                                    # mask only valid locations
print ('mask2 = ...')
print (mask2)

t[mask2] = mask_value
print ('t (after masking) = ...')
print (t)

print ('some consistency checks for mask2:')
print ('mask.sum (1):          ', mask.sum (1))
print ('mask2.sum (1):         ', mask2.sum (1))
print ('(mask2 * ~mask).any(): ', (mask2 * ~mask).any())
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.5.1
mask = ...
tensor([[False, False, False, False, False, False,  True,  True],
        [False, False, False, False,  True, False, False,  True],
        [False,  True,  True,  True,  True, False, False, False],
        [ True,  True, False,  True,  True, False,  True,  True],
        [ True, False, False,  True,  True,  True,  True,  True]])
mask2 = ...
tensor([[False, False, False, False, False, False, False,  True],
        [False, False, False, False,  True, False, False,  True],
        [False,  True,  True, False, False, False, False, False],
        [False, False, False,  True,  True, False,  True, False],
        [ True, False, False, False,  True,  True,  True,  True]])
t (after masking) = ...
tensor([[  0.,   0.,   0.,   0.,   0.,   0.,   0., 999.],
        [  0.,   0.,   0.,   0., 999.,   0.,   0., 999.],
        [  0., 999., 999.,   0.,   0.,   0.,   0.,   0.],
        [  0.,   0.,   0., 999., 999.,   0., 999.,   0.],
        [999.,   0.,   0.,   0., 999., 999., 999., 999.]])
some consistency checks for mask2:
mask.sum (1):           tensor([2, 2, 4, 6, 6])
mask2.sum (1):          tensor([1, 2, 2, 3, 5])
(mask2 * ~mask).any():  tensor(False)
</code></pre>
<p>*)  The proviso has to do with the probability of masking a specific<br>
number of elements in a given row.</p>
<p>For example, the index-2 row (the “third” row) of the <code>mask</code> tensor in<br>
my example script has four valid mask locations.  In your example<br>
code, the probability of masking three of those locations is the same<br>
as the probability of masking four of those locations (namely 25% in<br>
both cases), while in my version, the probability of masking three<br>
locations is higher that that of masking four locations.</p>
<p>It would depend on your use case as to whether this difference in<br>
probabilities matters.</p>
<p>(As an aside, I suspect that there is a way to get the same probability<br>
for masking any given number of valid mask locations, but I do not<br>
know how to do it.  My guess is that – if possible – it would take some<br>
careful and clever work with probability distributions to do it.)</p>
<p>Best.</p>
<p>K. Frank</p>