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()
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()
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)
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 !
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
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}})