[fx_acc] [fx2trt] add acc op mapper for argmin and converter for topk (#63823)
authorShiyan Deng <dsy842974287@fb.com>
Thu, 26 Aug 2021 20:06:46 +0000 (13:06 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Thu, 26 Aug 2021 20:16:22 +0000 (13:16 -0700)
commit7cfbc85821e8928db570a0730437b96484ac7b60
tree9cfb012d8ef538262dabefed4a1171e542b733e4
parentcbfec02007775d96139d8a1b9d9f8a44fcede31c
[fx_acc] [fx2trt] add acc op mapper for argmin and converter for topk (#63823)

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

Add mapper for `torch.argmin` which maps it to `acc_ops.flatten` (optional) + `acc_ops.topk` + `acc_ops.getitem` + `acc_ops.squeeze` (optional). This diff doesn't allow mapping if `dim=None && keepdim=True` in `torch.argmin`.

Add fx2trt converter for `acc_ops.topk`.

Test Plan:
buck test mode/opt glow/fb/fx/oss_acc_tracer:test_acc_tracer -- test_argmin
buck run mode/opt caffe2/torch/fb/fx2trt:test_topk

Reviewed By: jfix71

Differential Revision: D30501771

fbshipit-source-id: 0babc45e69bac5e61ff0b9b4dfb98940398e3e57
torch/fx/experimental/fx2trt/converters/acc_ops_converters.py
torch/fx/experimental/fx2trt/fx2trt.py
torch/fx/experimental/fx_acc/acc_ops.py