pytorch
987f1461 - [fx] Improve support for tuple subclasses such as NamedTuple (#73198)

Commit
2 years ago
[fx] Improve support for tuple subclasses such as NamedTuple (#73198) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/73198 Previously, if an arg to an FX node is a subclass of tuple then it gets sanitized essentially back to that base class. An example here is when setting an arg to be a TensorMetadata object, which is a NamedTuple, it will be set as a tuple instead. - Change `map_aggregate` to repack the tuple to `type(a)` when it's not directly a tuple (try/except for best attempt) - During codegen, call `add_global` for `type(a)` if it's not directly a tuple. - Add an option for an arg to provide a `_custom_fx_repr_fn` for use inside stringifying via `_format_arg` Test Plan: Added unit test coverage, where we inline the named tuple into arg/kwarg. Reviewed By: jamesr66a Differential Revision: D34381888 fbshipit-source-id: bd672a8542e2bba5aa604b448bec920efc256440 (cherry picked from commit 68f99c12dd40954b963490b90e821b315ba47048)
Author
Committer
Parents
Loading