[jit][script] enable advanced indexing with tensors #10862
Conversation
zou3519
requested review from
apaszke,
colesbury,
ezyang,
gchanan,
soumith and
zdevito
as code owners
jamesr66a
changed the title
torch/csrc/jit/script/compiler.cpp
Outdated
| auto handle_tensor = [&](Value* tensor) { | ||
| tensor_indices.resize(dim + 1); | ||
| tensor_indices[dim] = tensor; | ||
| dim++; |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
torch/csrc/jit/script/compiler.cpp
Outdated
| loc, gatherable.value(*graph), /*dim=*/0, index); | ||
| } | ||
| throw ErrorReport(loc) | ||
| << "Cannot index tensor with value of type " << index_type->str(); |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
zou3519
force-pushed
the
tensor-slicing
branch
from
55a1a73 to
6ff3d6c
Compare
zou3519
force-pushed
the
tensor-slicing
branch
from
6ff3d6c to
87b1613
Compare
There was a problem hiding this comment.
zou3519 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
There was a problem hiding this comment.
zou3519 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
zou3519
force-pushed
the
tensor-slicing
branch
from
7f1b22e to
125589f
Compare
There was a problem hiding this comment.
zou3519 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
On the way to pytorch#10774 This PR adds advanced indexing with tensors. The approach is to desugar advanced indexing into an at::index op. This is exactly how normal pytorch does it. [(I used this code as reference)](https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/python_variable_indexing.cpp) Sequences are a little tricky because JIT script doesn't support an easy way to turn arbitrary n-dimensional python lists into a tensor (it would be easy if we supported `torch.tensor`), so that'll come in a future PR.
- Add note about what the holes in tensor_indices are - Do some refactoring to make the code cleaner
zou3519
force-pushed
the
tensor-slicing
branch
from
125589f to
dfb3ed3
Compare
There was a problem hiding this comment.
zou3519 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
Summary: On the way to pytorch#10774 This PR adds advanced indexing with tensors. The approach is to desugar advanced indexing into an at::index op. This is exactly how normal pytorch does it. [(I used this code as reference)](https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/python_variable_indexing.cpp) Supporting sequences is a little tricky because JIT script doesn't have an easy way to turn arbitrary n-dimensional python lists into a tensor (it would be easy if we supported `torch.tensor`), so that'll come in a future PR. cc jamesr66a zdevito Pull Request resolved: pytorch#10862 Differential Revision: D9659449 Pulled By: zou3519 fbshipit-source-id: 56d293720d44c0fd27909e18327ab3985ddfced6
On the way to #10774
This PR adds advanced indexing with tensors.
The approach is to desugar advanced indexing into an at::index op.
This is exactly how normal pytorch does it.
(I used this code as reference)
Supporting sequences is a little tricky because JIT script doesn't have
an easy way to turn arbitrary n-dimensional python lists into a tensor
(it would be easy if we supported
torch.tensor), so that'll comein a future PR.
cc @jamesr66a @zdevito