[jit] Add start and step parameters for range in torchscript #20795
Conversation
facebook-github-bot
left a comment
facebook-github-bot
left a comment
There was a problem hiding this comment.
@Chillee has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
I decided to add special casing for the common case for a couple reasons:
|
facebook-github-bot
left a comment
facebook-github-bot
left a comment
There was a problem hiding this comment.
@Chillee has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
jamesr66a
left a comment
jamesr66a
left a comment
There was a problem hiding this comment.
Looking good! Some small stuff to fixup and we'll have to look at the tests but I think we're on the right track here
eellison
left a comment
eellison
left a comment
There was a problem hiding this comment.
Looks great 🔥a few comments
zdevito
left a comment
zdevito
left a comment
There was a problem hiding this comment.
Cool -- I think we should make the range/scale computations either TorchScript code or a builtin. Otherwise it will be too hard to maintain.
|
|
pytorchbot
added
the
module: internals
Chillee
force-pushed
the
jitrangeloops
branch
2 times, most recently
from
ba66992 to
9bed93a
Compare
Chillee
force-pushed
the
jitrangeloops
branch
2 times, most recently
from
54cfcf2 to
8ff19ea
Compare
jamesr66a
left a comment
jamesr66a
left a comment
There was a problem hiding this comment.
Looks good. I added some comments about small stuff
facebook-github-bot
left a comment
facebook-github-bot
left a comment
There was a problem hiding this comment.
@Chillee is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
@pytorchbot rebase this please. |
facebook-github-bot
left a comment
facebook-github-bot
left a comment
There was a problem hiding this comment.
@Chillee is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
Summary: Fixes #18440 I calculate a derived index from `start,stop,step` as `start + step*index`. When `start=0` and `step=1` (the defaults/`range(n)`), this is the same behavior as before. Unluckily, it seems that we do not optimize out operations like `x*1` or `x+0`. That means that we're doing lots of redundant operations when we don't need to. EDIT: More specifically, it seems like we only do this optimization for (tensor, scalar): https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/peephole.cpp#L128 The most annoying part of this code is calculating the number of iterations, given `start, stop, step`. I ended up going with the formula `(abs(stop-start) + abs(step)-1)//abs(step)`. Other intuitively appealing formulas like `(stop-start + step -1)//step` don't work for negative numbers. I tried using `SymbolicVariable` for the calculations, but it seems that `symbolicvariable` only outputs ops for `tensors`, not the integers we have. Pull Request resolved: pytorch/pytorch#20795 Differential Revision: D15446869 Pulled By: Chillee fbshipit-source-id: 6085545ace04e25985c6ac870226f7a651f670d5
8 participants
Fixes #18440
I calculate a derived index from
start,stop,stepasstart + step*index. Whenstart=0andstep=1(the defaults/range(n)), this is the same behavior as before.Unluckily, it seems that we do not optimize out operations like
x*1orx+0. That means that we're doing lots of redundant operations when we don't need to. EDIT: More specifically, it seems like we only do this optimization for (tensor, scalar): https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/peephole.cpp#L128The most annoying part of this code is calculating the number of iterations, given
start, stop, step. I ended up going with the formula(abs(stop-start) + abs(step)-1)//abs(step). Other intuitively appealing formulas like(stop-start + step -1)//stepdon't work for negative numbers.I tried using
SymbolicVariablefor the calculations, but it seems thatsymbolicvariableonly outputs ops fortensors, not the integers we have.