pytorch
194262ee - Make HigherOrderOperator stop appearing like torch.ops.* in FX (#103108)

Commit
1 year ago
Make HigherOrderOperator stop appearing like torch.ops.* in FX (#103108) Previously, defining a HigherOrderOperators (like cond) automatically generates a torch.ops.cond and causes them to trace into the FX graph as e.g. torch.ops.cond. This is not good, because: - Duplication. Since HigherOrderOperators are written in Python, they have an associated Python function that users should access them from. E.g. torch.cond (when we make it public). That is what should actually appear in the graph. - torch.ops.cond is a valid namespace for operator registration; having it be a function too confuses things. This PR: - Moves cond/map HigherOrderOperators to be under torch (necessary for the FX logic to not do weird things) - Sets the `__module__` of a HigherOrderOperator correct. This is what FX uses when tracing the operator. Test Plan: - updated tests Future: - I'll delete the ability to call cond as torch.ops.cond in a couple of days, after this change circulates internally. Pull Request resolved: https://github.com/pytorch/pytorch/pull/103108 Approved by: https://github.com/ydwu4
Author
Committer
Parents
Loading