[JIT] UseVariadicOp takes list_idx parameter (#63915)
authorMike Iovine <mikeiovine@fb.com>
Thu, 26 Aug 2021 21:09:10 +0000 (14:09 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Thu, 26 Aug 2021 21:10:35 +0000 (14:10 -0700)
commitaeec177833cb20e8c6177ef8dbcf02ddc37c8a32
tree3100d6159221b08ce6e230baac598465e75ace7d
parentd8d8e4902a9ce3426e84817b936699f85a5f698e
[JIT] UseVariadicOp takes list_idx parameter (#63915)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63915

Previously, this function only worked for variadic op substitutions of the form `op(list, args) -> variadic_op(list_1, ..., list_n, args)`. This change allows for transformations of the form `op(args_0, list, args_1) -> variadic_op(args_0, list_1, ..., list_n, args_1)`.

Test Plan:
`buck test caffe2/test/cpp/jit:jit -- Stack Concat`

(tests exercising `list_idx != 0` will be added further up in this diff stack)

Reviewed By: navahgar

Differential Revision: D30529729

fbshipit-source-id: 568080679c3b40bdaedee56bef2e8a5ce7985d2f
torch/csrc/jit/passes/variadic_ops.cpp
torch/csrc/jit/passes/variadic_ops.h