Indexing in C++

I’m trying to replicate the following Python code in the PyTorch C++ frontend:

value = torch::rand({2, 2, 4});
value = value[0, :1, ...]

I’m not sure how this’d look like if I’m using tensor.narrow() or tensor.select()

1 Like

The 0 in value[0, :1, ...] is equivalent to value.select(0, 0), and the :1 is equivalent to .narrow(0, 0, 1), then:

value.select(0, 0).narrow(0, 0, 1) == value[0, :1, ...]

Pay attention to the little catch: you can be tempted to do narrow(1, 0, 1) since you are narrowing the second dimension, but you will actually be narrowing the tensor after the selection, which will make first dimension disappear. What I mean is:

value.select(0, 0).narrow(0, 0, 1) == value.narrow(1, 0, 1).select(0, 0)
1 Like

Thanks for answering Levi! So to generalize things
value[a, :b, ...] is equivalent to value.select(a, 0).narrow(0, 0, b) and value[:b, a, ...] is equivalent to value.narrow(1, 0, b).select(a, 0)?

You can check it by yourself :smile: !

import torch
value = torch.rand(3, 4, 5, 6, 7)
a = 2
b = 1
value.select(0, a).narrow(0, 0, b).equal(value[a, :b, ...]) # True
value.narrow(1, 0, b).select(0, a).equal(value[a, :b, ...]) # True
2 Likes

Starting from the current nightly build (and PyTorch 1.5 soon), for

value = value[0, :1, ...]

we can write

using namespace torch::indexing;
value = value.index({0, Slice(None, 1), "..."});

Here is the general translation for Tensor::index and Tensor::index_put_ functions:

Python             C++ (assuming `using namespace torch::indexing`)
-------------------------------------------------------------------
0                  0
None               None
...                "..." or Ellipsis
:                  Slice()
start:stop:step    Slice(start, stop, step)
True / False       true / false
[[1, 2]]           torch::tensor({{1, 2}})
9 Likes