pytorch
1b147a52 - Allow FX tracer to trace control flow (if/while) statements when parameter shapes are in the conditionals (#61820)

Commit
3 years ago
Allow FX tracer to trace control flow (if/while) statements when parameter shapes are in the conditionals (#61820) Summary: Fixes https://github.com/pytorch/pytorch/issues/61733 Allow FX tracer to trace control flow (if/while) statements when parameter shapes are in the condition. If the user specifies the new "param_shapes_constant" option when constructing a tracer, the model's parameter shape attribute will be evaluated and the resulting constant will be emitted into the IR during tracing. Also added a new test ` python test/fx/test_fx_param_shape_control_flow.py ` The test also performs a somewhat whitebox style testing to check the generated Python code from the IR. Pull Request resolved: https://github.com/pytorch/pytorch/pull/61820 Reviewed By: bdhirsh Differential Revision: D29969299 Pulled By: jerryzhenleicai fbshipit-source-id: 99aae824bdfec880be69258de7ead5c8cd59eddc
Author
Jerry Cai
Parents
Loading