[dynamo] Be stricter about HigherOrderOperator kwargs (#111938)

jon-chuang · pytorchmergebot · commit edafe2ddb99d · 2023-10-26T03:51:30.000Z
kwargs need to be handled carefully in speculate subgraph. We should be clearer about the contract of what the inputs are. Pull Request resolved: #111938 Approved by: https://github.com/zou3519
diff --git a/torch/_dynamo/variables/higher_order_ops.py b/torch/_dynamo/variables/higher_order_ops.py
@@ -363,12 +363,6 @@ def make(value, source=None, **kwargs):
         else:
             unimplemented(f"HigherOrderOperator {value.__name__}")
 
-    def check_kwargs(self, kwargs, supported_types):
-        if not all(isinstance(value, supported_types) for value in kwargs.values()):
-            raise unimplemented(
-                f"Only kwargs of the following types are supported: {supported_types}"
-            )
-
     def call_function(
         self, tx, args: List[VariableTracker], kwargs: Dict[str, VariableTracker]
     ) -> VariableTracker:
@@ -616,6 +610,11 @@ def call_function(
         )
         from .builder import wrap_fx_proxy
 
+        if len(kwargs) > 0:
+            unsupported(
+                "torch.ops.higher_order.map: kwargs are not supported in the map operator."
+            )
+
         assert type(args[0]) in (UserFunctionVariable, NestedUserFunctionVariable)
         assert type(args[1]) is TensorVariable
 
@@ -672,16 +671,14 @@ def call_function(
         r = body_r.as_proxy().node.meta["example_value"]
         example_value = r.new_empty([sample_shape[0], *r.shape])
 
-        _, p_kwargs = proxy_args_kwargs([], kwargs)
-
         # Store the invocation as a call
         return wrap_fx_proxy(
             tx=tx,
             proxy=tx.output.create_proxy(
                 "call_function",
                 self.value,
                 args=tuple(p_args),
-                kwargs=p_kwargs,
+                kwargs={},
             ),
             example_value=example_value,
         )
@@ -691,17 +688,18 @@ class ExecutorchCallDelegateHigherOrderVariable(TorchHigherOrderOperatorVariable
     def call_function(
         self, tx, args: "List[VariableTracker]", kwargs: "Dict[str, VariableTracker]"
     ) -> "VariableTracker":
-        from . import ConstantVariable
         from .builder import wrap_fx_proxy
 
-        self.check_kwargs(kwargs, ConstantVariable)
-
         # This is operator for delegation within Executorch which calls a
         # specific function in the given lowered module with the given
         # operators. The actual operator is defined in the Executorch codebase.
         # This is a bad hierarchical violation since
         # executorch_call_delegate sits at a higher level than dynamo, but
         # there's no real solution to this issue yet.
+        if len(kwargs) > 0:
+            unimplemented(
+                "executorch_call_delegate: kwargs arguments were not enabled."
+            )
         lowered_module = tx.output.get_submodule(args[0].module_key)
 
         lowered_node = make_attr(tx, args[0].module_key)
@@ -715,8 +713,6 @@ def call_function(
 
         p_args = (lowered_node,) + p_args
 
-        _, p_kwargs = proxy_args_kwargs([], kwargs)
-
         # Store the invocation as a call
         return wrap_fx_proxy(
             tx=tx,
@@ -934,7 +930,7 @@ def call_function(
         if not isinstance(out_dims, (ConstantVariable, TupleVariable)):
             unimplemented("torch.func.vmap: out_dims is not an int or tuple variable.")
 
-        if kwargs:
+        if len(kwargs) > 0:
             unimplemented(
                 "NYI - torch.func.vmap: kwargs arguments are currently unsupported."
             )
@@ -1086,12 +1082,15 @@ def __init__(
     def call_function(
         self, tx, args: "List[VariableTracker]", kwargs: "Dict[str, VariableTracker]"
     ) -> "VariableTracker":
-        from . import ConstantVariable, UserFunctionVariable
+        from . import UserFunctionVariable
         from .builder import wrap_fx_proxy
 
         tracer = self.fwd_bwd_tracer
 
-        self.check_kwargs(kwargs, ConstantVariable)
+        if len(kwargs) > 0:
+            unimplemented(
+                "kwargs have not been implemented for torch.autograd.Function"
+            )
 
         from . import TorchVariable
 
@@ -1148,16 +1147,14 @@ def call_function(
         r = body_r.as_proxy().node.meta["example_value"]
         example_value = r
 
-        _, p_kwargs = proxy_args_kwargs([], kwargs)
-
         # Store the invocation as a call
         return wrap_fx_proxy(
             tx=tx,
             proxy=tx.output.create_proxy(
                 "call_function",
                 self.value,
                 args=tuple(p_args),
-                kwargs=p_kwargs,
+                kwargs={},
             ),
             example_value=example_value,
         )
@@ -1213,18 +1210,22 @@ def call_function(
     ) -> "VariableTracker":
         from .builder import wrap_fx_proxy
 
+        # This flattens the kwargs into lifted args
         p_args, p_kwargs, example_value, treespec = self.create_wrapped_node(
             tx, args, kwargs, "wrap"
         )
 
+        if len(p_kwargs) > 0:
+            unimplemented("kwargs should have been flattened into lifted args")
+
         # Store the invocation as a call
         variable = wrap_fx_proxy(
             tx=tx,
             proxy=tx.output.create_proxy(
                 "call_function",
                 self.value,
                 args=tuple(p_args),
-                kwargs=p_kwargs,
+                kwargs={},
             ),
             example_value=example_value,
         )
@@ -1246,7 +1247,7 @@ def call_function(
     ) -> "VariableTracker":
         from .builder import wrap_fx_proxy
 
-        if len(kwargs) != 0:
+        if len(kwargs) > 0:
             unimplemented("out_dtype does not handle kwargs")
 
         p_args = tuple(arg.as_proxy() for arg in args)
