mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
7691064768
Differential Revision: [D65307961](https://our.internmc.facebook.com/intern/diff/D65307961/) This PR introduces the concept of a "dispatcher" module `n` that carries multiple interpreter modules `n`, `n@1`, `n@2`, etc., each corresponding to a particular call of `n` and thus might carry a different specialized graph. We only do this when we're preserving module call signatures for `n`. The carried modules have the same number and order of calls to `n` appearing in the original module / exported program. In the unflattened module, all those calls go to the "dispatcher" module which internally tracks how many calls have been made so far and invokes the corresponding interpreter module. We reset this tracking after a successful or unsuccessful run of the unflattened module. Overall this makes swapping easier when module call signatures are preserved. Pull Request resolved: https://github.com/pytorch/pytorch/pull/139439 Approved by: https://github.com/tugsbayasgalan ghstack dependencies: #139438