4D boolean indexing with masks of lower dimmension

David_Juergens: <p>Hello,</p>
<p>I have a 4D tensor <code>A</code> which is boolean and of shape (L,n,L,a), and two 1D tensors <code>mask1</code> and <code>mask2</code> both of shape (L,).</p>
<p>I would like to be able to do the following:<br>
<code>A[mask1,0,mask2,:3] = True</code><br>
followed by<br>
<code>A[mask2,0,mask1,1:2] = True</code></p>
<p>Of course, the above does not work and results in shape mismatch. How can I accomplish this indexing/setting in a different way?</p>

KFrank: <p>Hi David!</p>
<p>Could you clarify your question a bit?</p>
<aside class="quote no-group" data-username="David_Juergens" data-post="1" data-topic="189010" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/david_juergens/48/53121_2.png" class="avatar"> David_Juergens:</div>
<blockquote>
<p>two 1D tensors <code>mask1</code> and <code>mask2</code> both of shape (L,).</p>
</blockquote>
</aside>
<p>What are the types of <code>mask1</code> and <code>mask2</code> and what would some example<br>
values for them be?</p>
<aside class="quote no-group">
<blockquote>
<p><code>A[mask1,0,mask2,:3] = True</code></p>
</blockquote>
</aside>
<p>As a simpler example, let’s say that <code>A</code> had shape <code>[L, L]</code>.  In such a case<br>
what, specifically, would you like <code>A[mask1, mask2] = True</code> to do?</p>
<p>Best.</p>
<p>K. Frank</p>

David_Juergens: <p>Hi K. Frank, thanks for your response.</p>
<p><code>mask1</code> and <code>mask2</code> are both boolean tensors filled with some amount of True and False along their 1 dimension.</p>
<p>In response to the second part of your question, in the end, I want <code>A</code> to actually be a mask for another tensor. Specifically, in my problem I have a loss function that outputs a (L,n,L,a) tensor of floats, but many of those values in the tensor I don’t need contributing to the final loss. In order to only have only specific entries contribute to the loss, I’m trying to create a mask tensor (<code>A</code>) that is also 4D, which only has <code>True</code> in the particular positions that I want contributing to the loss. I will then multiply the 4D loss tensor and the 4D mask together to zero out the unneeded contributions. Hopefully this clarifies a bit of context.</p>
<p>Here’s also a more concrete example:</p>
<pre><code class="lang-auto">import torch 

L = 100
n = 5
a = 13

A = torch.zeros((L,n,L,a)).bool()

mask1 = torch.randint(0,2,(L,)).bool()
mask2 = torch.randint(0,2,(L,)).bool()

print(A[mask1,0].shape) # works fine 

print(A[mask1,0,mask2,:3]) # error 

</code></pre>

KFrank: <p>Hi David!</p>
<aside class="quote no-group quote-modified" data-username="David_Juergens" data-post="3" data-topic="189010" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/david_juergens/48/53121_2.png" class="avatar"> David_Juergens:</div>
<blockquote>
<p><code>mask1</code> and <code>mask2</code> are both boolean tensors<br>
…</p>
<pre><code class="lang-auto">print(A[mask1,0,mask2,:3]) # error 
</code></pre>
</blockquote>
</aside>
<p>It’s still not clear precisely what you want to do.</p>
<p>Speaking in terms of the simplified two-dimensional example, it looks like<br>
you want to use <code>mask1</code> to mask the rows of <code>A</code> and use <code>mask2</code> to mask the<br>
columns.</p>
<p>But how should this work?</p>
<p>Do you want to keep <em>all</em> rows where <code>mask1</code> is <code>True</code> and <em>all</em> columns where<br>
<code>mask2</code> is <code>True</code>?</p>
<p>Do you only want to keep elements for which both <code>mask1</code> and <code>mask2</code> are<br>
<code>True</code>?</p>
<p>Or some other use case?</p>
<p>Here’s an illustration of these alternatives:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.1
&gt;&gt;&gt;
&gt;&gt;&gt; mask1 = torch.tensor ([True, True, False, False])
&gt;&gt;&gt; mask2 = torch.tensor ([True, False, True, False])
&gt;&gt;&gt;
&gt;&gt;&gt; t = torch.arange (16.0).reshape (4, 4) + 1.0
&gt;&gt;&gt;
&gt;&gt;&gt; t
tensor([[ 1.,  2.,  3.,  4.],
        [ 5.,  6.,  7.,  8.],
        [ 9., 10., 11., 12.],
        [13., 14., 15., 16.]])
&gt;&gt;&gt; t[torch.logical_and (mask1.unsqueeze (1), mask2.unsqueeze (0))]
tensor([1., 3., 5., 7.])
&gt;&gt;&gt; t[torch.logical_or (mask1.unsqueeze (1), mask2.unsqueeze (0))]
tensor([ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 11., 13., 15.])
&gt;&gt;&gt;
&gt;&gt;&gt; t[torch.logical_not (torch.logical_and (mask1.unsqueeze (1), mask2.unsqueeze (0)))] = 0.0
&gt;&gt;&gt; t
tensor([[1., 0., 3., 0.],
        [5., 0., 7., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>