Comprehensive documentation for Tensor indexing?

[gchaperon]

📚 The doc issue

Is there a comprehensive doc page for Tensor indexing, like there is for numpy.ndarray?

I am facing some inconsistencies, particularly when using combined advanced boolean indexing and basic indexing, where the behaviour is clearly documented in the numpy side of things, but I am getting IndexError s in pytorch.

Would be great to know to what extent Tensor indexing mimics ndarray indexing.

Cheers.

Suggest a potential alternative/fix

Add a page like https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing to the docs.

cc @svekars @holly1238

[ngimel]

Can you provide examples where behavior is divergent from numpy?

[gchaperon]

Sure.
I have a 3-dim tensor, and I want to index the first 2 using a 2-dim boolean mask, and index the third dimension using a slice object. Here is what I have:

t = torch.rand(3, 4, 10)
mask = torch.rand(3, 4) < 0.5

# raises IndexError: The shape of the mask [3, 4] at index 1 does not match the shape of the indexed tensor [3, 3, 10] at index 1 
t[mask, 1:]

The relevant numpy docs state that the above indexing should be equivalent to ...[mask.nonzero() + (slice(1, None, None),)] (using numpy's nonzero(), pytorch's requires an extra as_tuple=True flag). This is indeed the case:

# True
(t.numpy()[mask.numpy(), 1:] == t.numpy()[mask.numpy().nonzero() + (slice(1, None, None), )]).all()

Using the equivalent indexing object I can get the expected result in pytorch, so i guess the capability is there and is just the syntax that differs:

# tensor(True)
(t[mask.nonzero(as_tuple=True) + (slice(1, None, None),)] == torch.from_numpy(t.numpy()[mask.numpy(), 1:])).all()

I am using python version 3.10.5, torch version '1.11.0+cu113' and numpy version 1.22.4.

Thanks for the quick response!

[gchaperon]

oh, and btw I opend a docs issue because I didn't find any doc page documenting tensor indexing (specially advanced indexing). The closest I could find was in Tensor Views, in a note at the bottom, but since it doesn't really say "pytorch indexing has the exact same semantics as numpy indexing" I didn't think of this as a bug