Indexing operators

1

Hi guys.

Having an example tensor

a = torch.arange(81).reshape(3, 9, 3)

I’ve faced such an issue. selecting range indexing over [] and : notation works similarly for 2nd and first dimensions separately. e.g.

# both
a[:, 0, [0]]
a[:, 0, 0:1]
# give
tensor([[ 0],
        [27],
        [54]])

# both
a[:, 1:2, 0]
a[:, [1], 0]
# give
tensor([[ 3],
        [30],
        [57]])

So I actually thought these are really the same operator just in different notations

Still they produce different results once using both dimensions

a[:, [0], [0]]
tensor([[ 0],
        [27],
        [54]])

a[:, 0:1, [0]]
tensor([[[ 0]],

        [[27]],

        [[54]]]


a[:, 0:1, 0:1]
tensor([[[ 0]],

        [[27]],

        [[54]]]

Do I understand it wrong? Shouldn’t these operators do the same thing?

And why actually a[:, [0], [0]] has different shape compared to next two expressions

Or maybe is this some kind of a bug?

Thank you

2

Hi jv!

What you’re seeing here is the expected result of so-called advanced indexing.

The : in the first position is giving you the full length-three first dimension of your
original tensor.

The two [0]s in the second and third positions are two “index lists” that work together
to perform advanced indexing on the second and third dimensions jointly. The index
lists have shape [1] so they give you an advanced-indexing result from the second
and third dimensions together that has shape [1].

Together with the : this gives you a result with an overall shape of [3, 1].

I haven’t looked at the most recent pytorch documentation on advanced indexing, but
it used to be pretty sparse. My recollection is that it used to say “we do what numpy
does” and my recollection is that that is correct. So for details, I recommend taking a
look at numpy’s advanced-indexing documentation.

I’ve made an example to illustrate what is going on. (For simplicity, I made your initial
tensor a little smaller and got rid of its first dimension which just goes along for the ride
without really affecting anything.)

Here is the code:

import torch
print (torch.__version__)

import numpy as np
print (np.__version__)

t = torch.arange (15).reshape (5, 3)

print (t.shape)
print (t)

# use advanced indexing jointly on both dimensions
# each list has shape [1], so the result has shape [1]

print (t[[0], [0]].shape)   # advanced indexing selects a tensor of shape [1]
print (t[[0], [0]])

print (t[0, 0].shape)       # indexing with a python scalar kills the dimension
print (t[0, 0])             # and this "ordinary" indexing selects a single element

print (t[0:1, [0]].shape)   # ordinary slicing selects a first dimension of length 1
print (t[0:1, [0]])         # and (trivial) advanced indexing gives an additional dimension of shape [1]

# some more interesting examples of advanced indexing

print (t[[4, 1], [1, 2]].shape)
print (t[[4, 1], [1, 2]])

# index lists have shape [2, 4, 2]

i0 = [
    [[3, 2], [4, 0], [1, 1], [0, 2]],
    [[0, 0], [1, 2], [2, 3], [4, 4]]
]

i1 = [
    [[2, 1], [2, 0], [1, 1], [0, 2]],
    [[0, 0], [1, 2], [2, 2], [0, 2]]
]

# so advanced-indexing result has shape [2, 4, 2]

print (t[i0, i1].shape)
print (t[i0, i1])

# numpy example -- pytorch appears to follow numpy's advanced-indexing semantics

a = np.arange (15).reshape (5, 3)

print (a.shape)
print (a)

print (a[i0, i1].shape)
print (a[i0, i1])

And here is its output:

2.11.0+cu130
2.4.4
torch.Size([5, 3])
tensor([[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11],
        [12, 13, 14]])
torch.Size([1])
tensor([0])
torch.Size([])
tensor(0)
torch.Size([1, 1])
tensor([[0]])
torch.Size([2])
tensor([13,  5])
torch.Size([2, 4, 2])
tensor([[[11,  7],
         [14,  0],
         [ 4,  4],
         [ 0,  8]],

        [[ 0,  0],
         [ 4,  8],
         [ 8, 11],
         [12, 14]]])
(5, 3)
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]
 [12 13 14]]
(2, 4, 2)
[[[11  7]
  [14  0]
  [ 4  4]
  [ 0  8]]

 [[ 0  0]
  [ 4  8]
  [ 8 11]
  [12 14]]]

Best.

K. Frank

3

Hey. @KFrank.

Thanks for laying it down.

This has cracked me a little :slight_smile: . As I got it right

i0 = [
    [[3, 2], [4, 0], [1, 1], [0, 2]],
    [[0, 0], [1, 2], [2, 3], [4, 4]]
]

i1 = [
    [[2, 1], [2, 0], [1, 1], [0, 2]],
    [[0, 0], [1, 2], [2, 2], [0, 2]]
]

# so
t[i0, i1]
# it would be the same as literal
[
    [[t[3, 2], t[2, 1]], [..], [..], [..]],
    [[..], [..], [..], [..]]
]

in another word

i - stripping the dimension

i:j - dimension subset

[i,j,z]- kind of mapping the dimension

Still it must take me a while to go along with arbitrary combinations of all these together.

As well as to check could the advanced indexing be interrupted by others and having different shapes/different/nested dimensions etc. ))

a[:,[1,2,3], :7, [[4,5], [3,4], [0, 1], [8, 9]]]

Thank you

4

Hi jv!

Yes.

No, unless the different shapes for the index lists are broadcastable.

So, as written, this won’t work.

But:

a[:, [1, 2, 3], :7, [3, 2, 1]]                # works, same shape
a[:, [1, 2, 3], :7, [1]]                      # works, broadcastable
a[:, [1, 2, 3], :7, [[1, 2, 3], [3, 2, 1]]]   # works, broadcastable
a[:, [1, 2, 3], :7, [[1. 2], [3, 4]]          # fails, not broadcastable

Best.

K. Frank

5

Hi @KFrank

Ah, got it!

So I got now the name for this `broadcastable` for deeper workout. Thank you.
As I suppose this is something related the same to the following issue describes

a = torch.tensor([[1, 2, 3], [4, 5, 6]])
b = torch.tensor([2, 4])
a > b             # fails due to shapes
a > b[:, None]    # pass