Added prod op to FX2TRT (#72284)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72284
This update adds the prod op to the fx2trt tool which is used to create a TensorRT engine for a PyTorch model.
Test Plan:
A new unit test was added to test that the op was added to the acc tracer. This text can be run using the following command: buck test --debug //caffe2/test:test_fx_acc_tracer -- --exact 'caffe2/test:test_fx_acc_tracer - test_prod (fx_acc.test_acc_tracer.AccTracerTest)'
A new suite of unit tests were also added for the conversion to tensorRT and can be tested using the following command: buck test mode/dev-nosan //caffe2/test/fx2trt/converters:test_prod
Please note that unfortunately unlike other pytorch reduce ops such as sum, the pytorch prod function does not support reducing more than 1 dimension at a time (the dim arg cannot be a tuple, only a single int is acceptable for prod). Therefore prod cannot utilize all of the reduce_op code.
https://pxl.cl/1Xpn8
https://pxl.cl/1Xpn9
Reviewed By: 842974287
Differential Revision: D33875336
fbshipit-source-id: f9340db3685d681b1cf4ffc3b9fd25d16914e231
(cherry picked from commit cfe48d3737e0056409de5fdca61ada47c2cfc243)
