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
