frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add self to CODEOWNERS for fx/proxy.py; warn against adding new node arg types (#147031)

Browse Source 
Not sure if there's a better way

Pull Request resolved: https://github.com/pytorch/pytorch/pull/147031
Approved by: https://github.com/StrongerXi
ghstack dependencies: #147016, #147012, #147013
This commit is contained in:
rzou
2025-02-12 15:07:13 -08:00
committed by PyTorch MergeBot
parent 9a883007a2
commit 92d448ff62
2 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
CODEOWNERS
View File

@ -102,6 +102,9 @@ test/test_type_promotion.py @mruberry
test/functorch/test_ops.py @zou3519 @chillee @kshitij12345
test/functorch/test_vmap.py @zou3519 @chillee @kshitij12345
# This is the file where people can add new argument types to torch.fx.
torch/fx/proxy.py @zou3519
# HOPs
torch/_higher_order_ops/*.py @zou3519
torch/_dynamo/variables/higher_order_ops.py @zou3519

11
torch/fx/proxy.py
View File

@ -289,6 +289,17 @@ class TracerBase:
Can be override to support more trace-specific types.
"""
# IMPORTANT: Are you here because you are trying to proxy a new type into
# the graph? Please Please Please contact someone on the PyTorch Compiler team;
# the considerations are subtle.
#
# 1) When you add a new type, all of the downstream consumers and pass writers
# need to handle the new type. torch.fx is intended to be easy to write
# passes for, so we will push back against new types.
# 2) In torch.compile's IR, there are only specific operations that go
# into the graph. In particular, Tensor operations should go into the graph,
# but non-Tensor operations shouldn't. What that means is that constructors
# for new types *SHOULD NOT* become nodes in the FX graph.
if isinstance(a, Proxy):
return a.node # most common arg type goes first
elif hasattr(a, "__fx_create_arg__"):