Delay RNN graph break until module call (#186614)
Dynamo intentionally graph breaks on nn.RNN, nn.GRU, and nn.LSTM when allow_rnn is disabled. The old implementation raised the unsupported error while wrapping the recurrent module attribute. For LOAD_ATTR/LOAD_METHOD bytecode this goes through the generic step graph-break path, which can compile the prefix but cannot build a continuation for the call site, so the rest of the frame runs eagerly and post-RNN code is never compiled. Return a DelayGraphBreakVariable for unsupported recurrent modules instead. This keeps the attribute load traceable enough to reach the CALL instruction, where Dynamo's regular call graph-break handling can run the recurrent module eagerly and resume tracing the continuation. This matches the existing delayed break pattern used for disabled optimized modules and avoids changing the experimental allow_rnn=True behavior. Add regression coverage for RNN, GRU, and LSTM modules that verifies Dynamo compiles both the pre-recurrent and post-recurrent regions. Fixes #141184 Generated by my agent Test Plan: - python test/dynamo/test_modules.py -k test_rnn_graph_break_resumes_after_call - python test/export/test_export.py -k test_export_dynamo_config - python test/dynamo/test_modules.py -k unsupportedmodule - lintrunner -a Pull Request resolved: https://github.com/pytorch/pytorch/pull/186614 Approved by: https://github.com/mlazos
J
Jason Ansel committed
d5424d7a04b5017925d5cd8c841eb75fd36880db
Parent: 31f0501
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
on 9/30/2026, 10:30:33 PM