[Python] How does pytorch indexing work with lists of lists vs lists of lists of lists vs...

Following examples show how you can get two different answers depending on whether you have list of lists, or lists of list of lists, or 2d pytorch tensor. What is the logic used to go from the indexing array to one of the two outputs?

# create tensor
x = torch.arange(16).reshape(4, 4)
tensor([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11],
[12, 13, 14, 15]])

# index with two lists
x[
[0,1,],[2,3]
]
tensor([2, 7])

# index with list of two lists, i.e. index is 2d-array
x[
[[0,1,],[2,3]]
]
tensor([2, 7])

# index with list of list of two lists, i.e. index is 3d array
x[
[[[0,1,], [2,3,]]]
]
tensor([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7]],

[[ 8, 9, 10, 11],
[12, 13, 14, 15]]])

# index with tensor made from list of 2 lists, i.e. index is 2d array
x[
torch.tensor([[0,1], [2,3]])
]
tensor([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7]],

[[ 8, 9, 10, 11],
[12, 13, 14, 15]]])

Continue reading...