Does PyTorch follows NumPy's indexing?

ado_sar: <h2><a name="p-465721-example-1-1" class="anchor" href="#p-465721-example-1-1"></a>Example 1</h2>
<p>I was playing around with indexing a 3D tensor with a boolean mask and I noticed that <code>x[mask] != x[mask.nonzero()]</code> (but see the Note), in contrast with <a href="https://numpy.org/doc/stable/user/basics.indexing.html#boolean-array-indexing" rel="noopener nofollow ugc">NumPy where these two operations are equivalent</a> <code>x[mask] == x[mask.nonzero()]</code>.</p>
<pre data-code-wrap="python"><code class="lang-python">import numpy as np
import torch

torch.manual_seed(42)

tensor = torch.arange(36).reshape(3, 4, 3)
mask = torch.randint(2, (3, 4), dtype=torch.bool)

arr = tensor.clone().numpy()
mask_arr = mask.clone().numpy()

print(
        'mask and mask.nonzero are equivalent in torch:',
        torch.equal(tensor[mask], tensor[mask.nonzero()]),
        )
print(
        'mask and mask.nonzero are equivalent in numpy:',
        np.array_equal(arr[mask_arr], arr[mask_arr.nonzero()]),
        )

print('tensor[mask]:\n', tensor[mask])
print('arr[mask_arr]:\n', arr[mask_arr])
</code></pre>
<p>and the output is:</p>
<pre><code class="lang-auto">mask and mask.nonzero are equivalent in torch: False
mask and mask.nonzero are equivalent in numpy: True
tensor[mask]:
 tensor([[ 3,  4,  5],
        [15, 16, 17],
        [27, 28, 29]])
arr[mask_arr]:
 [[ 3  4  5]
 [15 16 17]
 [27 28 29]]
</code></pre>
<h3><a name="p-465721-note-2" class="anchor" href="#p-465721-note-2"></a>Note</h3>
<p>According to the <a href="https://pytorch.org/docs/main/generated/torch.nonzero.html" rel="noopener nofollow ugc">PyTorch docs</a>:</p>
<blockquote>
<p><code>torch.nonzero(..., as_tuple=True)</code> returns a tuple of 1-D index tensors, allowing for advanced indexing …</p>
</blockquote>
<p>If we use <code>x[mask.nonzero(as_tuple=True)]</code>, then the output is equivalent to <code>x[mask]</code>, so I assume that happens under the hood when using boolean tensors for advanced indexing. Is that correct?</p>
<h2><a name="p-465721-example-2-3" class="anchor" href="#p-465721-example-2-3"></a>Example 2</h2>
<p>Another example where PyTorch and NumPy differ is the following:</p>
<pre><code class="lang-auto">import torch
import numpy as np

a = torch.arange(9).reshape(3, 3)
b = np.arange(9).reshape(3, 3)

print('x[[[0]]] in torch:\n', a[[[0]]])
print('x[[[0]]] in numpy:\n', b[[[0]]])
</code></pre>
<p>and as shown in the output the shapes of the resulting matrices are different:</p>
<pre><code class="lang-auto">x[[[0]]] in torch:
 tensor([[0, 1, 2]])
x[[[0]]] in numpy:
 [[[0 1 2]]]
</code></pre>
<h2><a name="p-465721-example-3-4" class="anchor" href="#p-465721-example-3-4"></a>Example 3</h2>
<p>This is related to: <a href="https://github.com/pytorch/pytorch/issues/6773" class="inline-onebox" rel="noopener nofollow ugc">Indexing with list of booleans differs from numpy · Issue #6773 · pytorch/pytorch · GitHub</a></p>
<pre data-code-wrap="python"><code class="lang-python">import numpy as np
import torch

np.random.seed(43)

a = np.arange(27).reshape(3, 3, 3)
b = np.random.randint(2, size=(3, 3)).astype('bool')
print(a[b.tolist()])

c = torch.tensor(a)
d = torch.tensor(b)
print(c[d.tolist()])
</code></pre>
<p>and the output is:</p>
<pre><code class="lang-auto">[[ 6  7  8]
 [ 9 10 11]
 [12 13 14]
 [21 22 23]
 [24 25 26]]
tensor([19, 23])
</code></pre>
<p>Are there any known cases where PyTorch and NumPy differ with regards to indexing?</p>