[dynamo] Fix dynamic shapes handling in after_aot repro generation

[skarjala]

Stack from ghstack (oldest at bottom):

• [dynamo] Move skipIf decorator to class level in test_fx_graph_runnable #157594
• [dynamo] Add FakeProcessGroup support for fx_graph_runnable with distributed collectives #157162
• -> [dynamo] Fix dynamic shapes handling in after_aot repro generation #157136

Summary:

• Extract symbolic variables directly from graph placeholders and arguments
• Add symbolic variable definitions to generated repro code
• Add unit tests with ToyModel for testing

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/157136

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (2 Unrelated Failures)

As of commit 067b745 with merge base 6c24c66 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh) (disabled by #148644 but the issue was closed recently and a rebase is needed to make it pass)
  dynamo/test_decorators.py::DecoratorTests::test_set_stance_aot_eager_then_compile

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

• pull / cuda12.8-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu, unstable) (gh) (#153987)
  MISSING REGRESSION TEST

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[skarjala]

@pytorchbot merge -i

[pytorchmergebot]

Merge started

Your change will be merged while ignoring the following 1 checks: pull / cuda12.8-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu, unstable)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 1 jobs have failed, first few of them are: linux-binary-manywheel / manywheel-py3_9-cuda12_9-build / build

Details for Dev Infra team

Raised by workflow job

[skarjala]

@pytorchbot merge -i

[pytorchmergebot]

Merge started

Your change will be merged while ignoring the following 2 checks: pull / cuda12.8-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu, unstable), trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here