pytorch
d52f8c25 - Add should_traverse_fn to torch.fx.node.map_aggregate (#81510)

Commit

2 years ago

Add should_traverse_fn to torch.fx.node.map_aggregate (#81510) Adds an optional callback that checks if map_aggregate should continue recursive traversal. The main motivation is to not traverse torch.Size which is tuple Pull Request resolved: https://github.com/pytorch/pytorch/pull/81510 Approved by: https://github.com/SherlockNoMad, https://github.com/jamesr66a

Author

pbelevich

Committer

pytorchmergebot

Parents

7af0200a

Files3

test
- expect
  - TestFXAPIBackwardCompatibility.test_function_back_compat-fx_backcompat_function_signatures.expect
- test_fx.py
torch/fx
- node.py

pytorch d52f8c25 - Add should_traverse_fn to torch.fx.node.map_aggregate (#81510)

pytorch
d52f8c25 - Add should_traverse_fn to torch.fx.node.map_aggregate (#81510)