Update how we input kwargs (#314)

Summary:

Previously, the code for passing inputs to exported program was:
```
if kwargs:
    return (args, kwargs)
else:
    return args
```

However, this causes some inconsistency where if the original input contains args and kwargs, the treespec would be a tuple containing a tuple of arguments, and a dictionary of keyword arguments. But if the original input only contained args, the treespec would just be a tuple of arguments.

So I updated the code to just always keep the kwargs around

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng

X-link: pytorch/pytorch#109160

Reviewed By: zhxchen17

Differential Revision: D49218534

Pulled By: angelayi

Author: angelayi

exir/capture/_capture.py

@@ -163,7 +163,7 @@ def capture(  # noqa: C901
         # input spec, but due to some limitations in pytree implementation, it doesn't
         # recognize the make_fx graph with torchdynamo input spec. We workaround it
         # by getting the input spec directly from user argument.
-        in_spec = pytree.tree_flatten(args)[1]
+        in_spec = pytree.tree_flatten((args, {}))[1]
 
         if config.enable_functionalization and not config.enable_aot:
             args = copy.deepcopy(args)

exir/emit/test/test_emit.py

@@ -173,7 +173,8 @@ def f(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]:
         )
 
         self.assertEqual(
-            program.execution_plan[0].container_meta_type.encoded_inp_str, "T1#1($)"
+            program.execution_plan[0].container_meta_type.encoded_inp_str,
+            "T2#1#0(T1#1($),D0())",
         )
 
     def test_buffers_with_perfect_alignment(self) -> None:

exir/program/_program.py

@@ -75,7 +75,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         from torch._export import combine_args_kwargs
 
         if self.call_spec.in_spec is not None:
-            user_args = combine_args_kwargs(args, kwargs)
+            user_args = args
             try:
                 args = fx_pytree.tree_flatten_spec(user_args, self.call_spec.in_spec)  # type: ignore[assignment]
             except Exception:

exir/tests/fixtures/basic_sin_max.txt

@@ -6,7 +6,7 @@
     {
       "name": "forward",
       "container_meta_type": {
-        "encoded_inp_str": "T1#1($)",
+        "encoded_inp_str": "T2#1#0(T1#1($),D0())",
         "encoded_out_str": "$"
       },
       "values": [

exir/tests/fixtures/composite_delegate.txt

@@ -6,7 +6,7 @@
     {
       "name": "forward",
       "container_meta_type": {
-        "encoded_inp_str": "T3#1#1#1($,$,$)",
+        "encoded_inp_str": "T2#3#0(T3#1#1#1($,$,$),D0())",
         "encoded_out_str": "$"
       },
       "values": [

exir/verification/interpreter.py

@@ -369,7 +369,7 @@ def run(self, *raw_args: torch.Tensor) -> PyTree:
         """
 
         # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
-        args, pytree = ex_pytree.tree_flatten(raw_args)
+        args, pytree = ex_pytree.tree_flatten((raw_args, {}))
 
         if pytree.to_str() != self.container_metatype.encoded_inp_str:
             raise TypeError(