Add self to CODEOWNERS for fx/proxy.py; warn against adding new node …

…arg types (pytorch#147031)

Not sure if there's a better way

Pull Request resolved: pytorch#147031
Approved by: https://github.com/StrongerXi
ghstack dependencies: pytorch#147016, pytorch#147012, pytorch#147013

CODEOWNERS

@@ -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

torch/fx/proxy.py

@@ -289,6 +289,17 @@ def create_arg(self, a: Any) -> Argument:
 
         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__"):