[fx][normalize] Allow for args to be left as args (pytorch#55995)

Summary:
Pull Request resolved: pytorch#55995

Normalization is kind of broken currently. But making default arguments visible still appears to work, and is nice functionality to still be able to rely on/use. Adds an option to `NormalizeArgs`'s `__init__` called `normalize_to_only_use_kwargs` which defaults to true, which if set to false will keep using the same signature as provided, but additionally set kwargs in kwargs.

Test Plan: Added test to `test_fx_experimental`.

Reviewed By: 842974287

Differential Revision: D27759448

fbshipit-source-id: 620061fcf46d8549ac70b62aede8b6740aee3778

Author: jfix71

test/test_fx_experimental.py

@@ -911,8 +911,8 @@ def forward(self, {params}):
                         normalized_args2 = normalize_module(traced, node.target, node.args, node.kwargs)
                         assert(normalized_args == normalized_args2)
                         assert normalized_args
-                        node.args = ()
-                        node.kwargs = normalized_args
+                        node.args = normalized_args.args
+                        node.kwargs = normalized_args.kwargs
 
             traced.recompile()
 
@@ -1306,8 +1306,8 @@ def jit_infer_type(v):
                 continue
             # Test normalize_function by itself
             ref_out = op.op(*arg_values, **kwarg_values)
-            norm_kwargs = normalize_function(op.op, arg_values, kwarg_values, arg_types, kwarg_types)
-            test_out = op.op(**norm_kwargs)
+            norm_args_and_kwargs = normalize_function(op.op, arg_values, kwarg_values, arg_types, kwarg_types)
+            test_out = op.op(*norm_args_and_kwargs.args, **norm_args_and_kwargs.kwargs)
             self.assertEqual(test_out, ref_out)
 
             # Test normalized_arguments as part of FX
@@ -1351,8 +1351,8 @@ def forward(self, {', '.join(param_names)}):
                 if node.op == 'call_function':
                     normalized_args = node.normalized_arguments(traced, arg_types, kwarg_types)
                     assert normalized_args
-                    node.args = ()
-                    node.kwargs = normalized_args
+                    node.args = normalized_args.args
+                    node.kwargs = normalized_args.kwargs
             traced.recompile()
 
             test_out = traced(*param_values)

torch/fx/experimental/normalize.py

@@ -13,9 +13,10 @@ class NormalizeArgs(Transformer):
     """
     Normalize arguments to Python targets. This means that
     `args/kwargs` will be matched up to the module/functional's
-    signature and rewritten to exclusively kwargs in positional order.
-    Also populates default values. Does not support positional-only
-    parameters or varargs parameters (*args, **kwargs).
+    signature and rewritten to exclusively kwargs in positional order
+    if `normalize_to_only_use_kwargs` is true. Also populates default
+    values. Does not support positional-only parameters or varargs
+    parameters (*args, **kwargs).
 
     If the nodes have 'type' metadata, it will use it to disambiguate
     overloads. Otherwise, it will throw an error.
@@ -25,9 +26,11 @@ class NormalizeArgs(Transformer):
         traced = torch.fx.symbolic_trace(m)
         traced = NormalizeArgs(traced).transform()
     """
-    def __init__(self, module : torch.nn.Module):
+    def __init__(self, module : torch.nn.Module,
+                 normalize_to_only_use_kwargs : bool = True):
         super().__init__(module)
         self.node_map: Dict[Proxy, Node] = {}
+        self.normalize_to_only_use_kwargs = normalize_to_only_use_kwargs
 
     def run_node(self, n: Node) -> Any:
         args, kwargs = self.fetch_args_kwargs_from_env(n)
@@ -51,17 +54,21 @@ def call_function(
             self, target : Target, args : Tuple[Argument, ...], kwargs : Dict[str, Any],
             arg_types: Optional[Tuple[Any, ...]] = None, kwarg_types : Optional[Dict[str, Any]] = None):
         assert callable(target)
-        new_kwargs = normalize_function(target, args, kwargs, arg_types, kwarg_types)  # type: ignore[arg-type]
-        if new_kwargs:
-            return self.tracer.create_proxy('call_function', target, (), new_kwargs)
+        new_args_and_kwargs = normalize_function(target, args, kwargs, arg_types, kwarg_types,  # type: ignore[arg-type]
+                                                 self.normalize_to_only_use_kwargs)
+        if new_args_and_kwargs:
+            new_args, new_kwargs = new_args_and_kwargs
+            return self.tracer.create_proxy('call_function', target, new_args, new_kwargs)
         else:
             return super().call_function(target, args, kwargs)
 
     def call_module(self, target : Target, args : Tuple[Argument, ...], kwargs : Dict[str, Any]):
         assert isinstance(target, str)
-        new_kwargs = normalize_module(self.module, target, args, kwargs)  # type: ignore[arg-type]
-        if new_kwargs:
-            return super().call_module(target, (), new_kwargs)
+        new_args_and_kwargs = normalize_module(self.module, target, args, kwargs,  # type: ignore[arg-type]
+                                               self.normalize_to_only_use_kwargs)
+        if new_args_and_kwargs:
+            new_args, new_kwargs = new_args_and_kwargs
+            return super().call_module(target, new_args, new_kwargs)
         else:
             return super().call_module(target, args, kwargs)
 

torch/fx/node.py

@@ -4,7 +4,7 @@
 import torch
 import builtins
 import types
-from torch.fx.operator_schemas import normalize_function, normalize_module
+from torch.fx.operator_schemas import normalize_function, normalize_module, ArgsKwargsPair
 
 if TYPE_CHECKING:
     from .graph import Graph
@@ -446,11 +446,13 @@ def is_impure(self):
 
     def normalized_arguments(
             self, root : torch.nn.Module, arg_types : Optional[Tuple[Any]] = None,
-            kwarg_types : Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
+            kwarg_types : Optional[Dict[str, Any]] = None,
+            normalize_to_only_use_kwargs : bool = False) -> Optional[ArgsKwargsPair]:
         """
         Returns normalized arguments to Python targets. This means that
         `args/kwargs` will be matched up to the module/functional's
-        signature and return exclusively kwargs in positional order.
+        signature and return exclusively kwargs in positional order
+        if `normalize_to_only_use_kwargs` is true.
         Also populates default values. Does not support positional-only
         parameters or varargs parameters.
 
@@ -462,10 +464,11 @@ def normalized_arguments(
             root (torch.nn.Module): Module upon which to resolve module targets.
             arg_types (Optional[Tuple[Any]]): Tuple of arg types for the args
             kwarg_types (Optional[Dict[str, Any]]): Dict of arg types for the kwargs
+            normalize_to_only_use_kwargs (bool): Whether to normalize to only use kwargs.
 
         Returns:
 
-            Returns normalized_kwargs, or `None` if not successful.
+            Returns NamedTuple ArgsKwargsPair, or `None` if not successful.
         """
         if self.op == 'call_function':
             assert callable(self.target)

torch/fx/operator_schemas.py

@@ -3,9 +3,16 @@
 import numbers
 import typing
 import enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, cast
+from typing import Any, Callable, Dict, List, Optional, Tuple, NamedTuple, cast
 from torch._jit_internal import boolean_dispatched
 
+class ArgsKwargsPair(NamedTuple):
+    """
+    Simple named tuple for wrapping args/kwargs pairs.
+    """
+    args: Tuple[Any, ...]
+    kwargs: Dict[str, Any]
+
 _manual_overrides : Dict[Callable, List[inspect.Signature]] = {}
 
 def _nonzero_schemas():
@@ -140,11 +147,13 @@ def is_homogeneous_int_tuple(t):
 
 def normalize_function(
         target: Callable, args: Tuple[Any], kwargs : Optional[Dict[str, Any]] = None, arg_types : Optional[Tuple[Any]] = None,
-        kwarg_types : Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
+        kwarg_types : Optional[Dict[str, Any]] = None,
+        normalize_to_only_use_kwargs : bool = False) -> Optional[ArgsKwargsPair]:
     """
     Returns normalized arguments to PyTorch functions. This means that
     `args/kwargs` will be matched up to the functional's
-    signature and return exclusively kwargs in positional order.
+    signature and return exclusively kwargs in positional order if
+    `normalize_to_only_use_kwargs` is True.
     Also populates default values. Does not support positional-only
     parameters or varargs parameters (*args, **kwargs). Does not support modules.
 
@@ -156,14 +165,15 @@ def normalize_function(
         kwargs (Optional[Dict[str, Any]]): Dict of kwargs to the function
         arg_types (Optional[Tuple[Any]]): Tuple of arg types for the args
         kwarg_types (Optional[Dict[str, Any]]): Dict of arg types for the kwargs
+        normalize_to_only_use_kwargs (bool): Whether to normalize to only use kwargs.
 
     Returns:
 
-        Returns normalized_kwargs, or `None` if not successful.
+        Returns normalized_args_and_kwargs, or `None` if not successful.
     """
     if kwargs is None:
         kwargs = {}
-    new_kwargs = None
+    new_args_and_kwargs = None
     if target in boolean_dispatched or target.__module__ in ['torch.nn.functional', 'torch.functional']:
         target_for_analysis = target
         if target in boolean_dispatched:
@@ -180,15 +190,15 @@ def normalize_function(
 
         assert callable(target_for_analysis)
         sig = inspect.signature(inspect.unwrap(target_for_analysis))
-        new_kwargs = _args_kwargs_to_normalized_kwargs(sig, args, kwargs)
+        new_args_and_kwargs = _args_kwargs_to_normalized_args_kwargs(sig, args, kwargs, normalize_to_only_use_kwargs)
     else:
         assert callable(target)
         torch_op_schemas = get_signature_for_torch_op(target)
         matched_schemas = []
         if torch_op_schemas:
             # Iterate through all of the schema until we find one that matches
-            # If one matches, populate `new_kwargs` with the combined args/kwargs
-            # values. If none matches, `new_kwargs` will be None
+            # If one matches, populate `new_args_and_kwargs` with the new args/kwargs
+            # values. If none matches, `new_args_and_kwargs` will be None
             for candidate_signature in torch_op_schemas:
                 try:
                     candidate_signature.bind(*args, **kwargs)
@@ -201,7 +211,8 @@ def normalize_function(
                 pass
             elif len(matched_schemas) == 1:
                 # Matched exactly one schema, unambiguous
-                new_kwargs = _args_kwargs_to_normalized_kwargs(matched_schemas[0], args, kwargs)
+                new_args_and_kwargs = _args_kwargs_to_normalized_args_kwargs(matched_schemas[0], args, kwargs,
+                                                                             normalize_to_only_use_kwargs)
             else:
                 if arg_types is not None or kwarg_types is not None:
                     arg_types = arg_types if arg_types else cast(Tuple[Any], ())
@@ -216,7 +227,8 @@ def normalize_function(
                         except TypeError as e:
                             sig_matches = False
                         if sig_matches:
-                            new_kwargs = _args_kwargs_to_normalized_kwargs(candidate_signature, args, kwargs)
+                            new_args_and_kwargs = _args_kwargs_to_normalized_args_kwargs(candidate_signature, args, kwargs,
+                                                                                         normalize_to_only_use_kwargs)
                             break
                 else:
                     # Matched more than one schema. In this situation, the caller must provide the types of
@@ -226,14 +238,16 @@ def normalize_function(
                                        f'the schema match was ambiguous! Please provide argument types to '
                                        f'the normalize_arguments() call. Available schemas:\n{schema_printouts}')
 
-    return new_kwargs
+    return new_args_and_kwargs
 
 def normalize_module(
-        root: torch.nn.Module, target: str, args: Tuple[Any], kwargs : Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
+        root: torch.nn.Module, target: str, args: Tuple[Any], kwargs : Optional[Dict[str, Any]] = None,
+        normalize_to_only_use_kwargs : bool = False) -> Optional[ArgsKwargsPair]:
     """
     Returns normalized arguments to PyTorch modules. This means that
     `args/kwargs` will be matched up to the functional's
-    signature and return exclusively kwargs in positional order.
+    signature and return exclusively kwargs in positional order if
+    `normalize_to_only_use_kwargs` is True.
     Also populates default values. Does not support positional-only
     parameters or varargs parameters (*args, **kwargs).
 
@@ -242,10 +256,11 @@ def normalize_module(
         target (Callable): Function that we are normalizing
         args (Tuple[Any]): Tuple of args to the function
         kwargs (Optional[Dict[str, Any]]): Dict of kwargs to the function
+        normalize_to_only_use_kwargs (bool): Whether to normalize to only use kwargs.
 
     Returns:
 
-        Returns normalized_kwargs, or `None` if not successful.
+        Returns normalized_args_and_kwargs, or `None` if not successful.
     """
     try:
         submod = root.get_submodule(target)
@@ -258,27 +273,30 @@ def normalize_module(
             sig = inspect.signature(inspect.unwrap(submod.forward))
             if kwargs is None:
                 kwargs = {}
-            new_kwargs = _args_kwargs_to_normalized_kwargs(sig, args, kwargs)
-            return new_kwargs
+            new_args_and_kwargs = _args_kwargs_to_normalized_args_kwargs(sig, args, kwargs,
+                                                                         normalize_to_only_use_kwargs)
+            return new_args_and_kwargs
     return None
 
-def _args_kwargs_to_normalized_kwargs(sig : inspect.Signature, args : Tuple[Any, ...],
-                                      kwargs : Dict[str, Any]) -> Optional[Dict[str, Any]]:
+def _args_kwargs_to_normalized_args_kwargs(sig : inspect.Signature, args : Tuple[Any, ...],
+                                           kwargs : Dict[str, Any],
+                                           normalize_to_only_use_kwargs : bool) -> Optional[ArgsKwargsPair]:
     """
     Given a call target, args, and kwargs, return the arguments normalized into
-    a single kwargs dict, or None if the type signature is not supported by
+    an ArgsKwargsPair, or None if the type signature is not supported by
     this normalization.
 
     Args:
 
         target (inspect.Signature): Signature object for the target
         args (Tuple): Arguments that appear at the callsite for `target`
         kwargs (Dict): Keyword arugments that appear at the callsite for `target`
+        normalize_to_only_use_kwargs (bool): Whether to normalize to only use kwargs.
 
     Returns:
 
-        Optional[Dict]: Normalized kwargs for `target`, or `None` if this target is not
-            supported
+        Optional[ArgsKwargsPair]: Normalized args and kwargs for `target`, or `None` if
+            this target is not supported.
     """
 
     # Don't currently support positional-only
@@ -292,7 +310,11 @@ def _args_kwargs_to_normalized_kwargs(sig : inspect.Signature, args : Tuple[Any,
     bound_args.apply_defaults()
 
     new_kwargs : Dict[str, Any] = {}
-    for param in sig.parameters:
-        new_kwargs[param] = bound_args.arguments[param]
-
-    return new_kwargs
+    new_args : List[Any] = []
+    for i, param in enumerate(sig.parameters):
+        if not normalize_to_only_use_kwargs and i < len(args):
+            new_args.append(bound_args.arguments[param])
+        else:
+            new_kwargs[param] = bound_args.arguments[param]
+
+    return ArgsKwargsPair(tuple(new_args), new_kwargs)