[export] Fix tree spec matching behavior.

test/export/test_export.py

@@ -1145,5 +1145,16 @@ def forward(self, x: Input):
                     else:
                         self.assertIsInstance(s, int)
 
+    def test_export_with_wrong_inputs(self):
+        class MyModule(torch.nn.Module):
+            def forward(self, x):
+                return x + x
+
+        exported_program = export(MyModule(), (torch.rand(2, 3),), {})
+        with self.assertRaisesRegex(
+            TypeError, "Trying to flatten user inputs with exported input tree spec"
+        ):
+            exported_program(torch.rand(2, 3), torch.rand(2, 3))
+
 if __name__ == '__main__':
     run_tests()

torch/export/exported_program.py

@@ -346,10 +346,12 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         if self.call_spec.in_spec is not None:
             try:
                 user_args = combine_args_kwargs(args, kwargs)
-                args = fx_pytree.tree_flatten_spec(user_args, self.call_spec.in_spec)  # type: ignore[assignment]
+                args = fx_pytree.tree_flatten_spec(
+                    user_args, self.call_spec.in_spec, exact_structural_match=True
+                )  # type: ignore[assignment]
             except Exception:
                 _, received_spec = pytree.tree_flatten(user_args)
-                raise error.InternalError(
+                raise TypeError(
                     "Trying to flatten user inputs with exported input tree spec: \n"
                     f"{self.call_spec.in_spec}\n"
                     "but actually got inputs with tree spec of: \n"

torch/fx/_pytree.py

@@ -1,14 +1,25 @@
-from typing import Callable, Any, Tuple, List, Dict, Type, NamedTuple
-from torch.utils._pytree import PyTree, TreeSpec, LeafSpec
 from collections import namedtuple
+from typing import Any, Callable, Dict, List, NamedTuple, Tuple, Type, Optional
+
+from torch.utils._pytree import LeafSpec, PyTree, TreeSpec
 
 FlattenFuncSpec = Callable[[PyTree, TreeSpec], List]
 
-SUPPORTED_NODES: Dict[Type[Any], Any] = {}
-def register_pytree_flatten_spec(typ: Any, flatten_fn_spec: FlattenFuncSpec) -> None:
+FlattenFuncExactMatchSpec = Callable[[PyTree, TreeSpec], bool]
+
+SUPPORTED_NODES: Dict[Type[Any], FlattenFuncSpec] = {}
+
+SUPPORTED_NODES_EXACT_MATCH: Dict[Type[Any], Optional[FlattenFuncExactMatchSpec]] = {}
+
+def register_pytree_flatten_spec(
+    typ: Any,
+    flatten_fn_spec: FlattenFuncSpec,
+    flatten_fn_exact_match_spec: Optional[FlattenFuncExactMatchSpec] = None
+) -> None:
     SUPPORTED_NODES[typ] = flatten_fn_spec
+    SUPPORTED_NODES_EXACT_MATCH[typ] = flatten_fn_exact_match_spec
 
-def tree_flatten_spec(pytree: PyTree, spec: TreeSpec) -> List[Any]:
+def tree_flatten_spec(pytree: PyTree, spec: TreeSpec, exact_structural_match=False) -> List[Any]:
     if isinstance(spec, LeafSpec):
         return [pytree]
     if spec.type not in SUPPORTED_NODES:
@@ -18,9 +29,13 @@ def tree_flatten_spec(pytree: PyTree, spec: TreeSpec) -> List[Any]:
             "sure that any custom pytrees have been registered before loading it.")
     flatten_fn_spec = SUPPORTED_NODES[spec.type]
     child_pytrees = flatten_fn_spec(pytree, spec)
+    if exact_structural_match:
+        flatten_fn_exact_match_spec = SUPPORTED_NODES_EXACT_MATCH[spec.type]
+        if flatten_fn_exact_match_spec and not flatten_fn_exact_match_spec(pytree, spec):
+            raise RuntimeError(f"Cannot flatten pytree {pytree}, given spec: {spec}")
     result = []
     for child, child_spec in zip(child_pytrees, spec.children_specs):
-        flat = tree_flatten_spec(child, child_spec)
+        flat = tree_flatten_spec(child, child_spec, exact_structural_match)
         result += flat
     return result
 
@@ -36,7 +51,19 @@ def _tuple_flatten_spec(d: Tuple[Any], spec: TreeSpec) -> List[Any]:
 def _namedtuple_flatten_spec(d: NamedTuple, spec: TreeSpec) -> List[Any]:
     return [d[i] for i in range(len(spec.children_specs))]
 
-register_pytree_flatten_spec(dict, _dict_flatten_spec)
-register_pytree_flatten_spec(list, _list_flatten_spec)
-register_pytree_flatten_spec(tuple, _tuple_flatten_spec)
-register_pytree_flatten_spec(namedtuple, _tuple_flatten_spec)
+def _dict_flatten_spec_exact_match(d: Dict[Any, Any], spec: TreeSpec) -> bool:
+    return len(d) == len(spec.context)
+
+def _list_flatten_spec_exact_match(d: List[Any], spec: TreeSpec) -> bool:
+    return len(d) == len(spec.children_specs)
+
+def _tuple_flatten_spec_exact_match(d: Tuple[Any], spec: TreeSpec) -> bool:
+    return len(d) == len(spec.children_specs)
+
+def _namedtuple_flatten_spec_exact_match(d: NamedTuple, spec: TreeSpec) -> bool:
+    return len(d) == len(spec.children_specs)
+
+register_pytree_flatten_spec(dict, _dict_flatten_spec, _dict_flatten_spec_exact_match)
+register_pytree_flatten_spec(list, _list_flatten_spec, _list_flatten_spec_exact_match)
+register_pytree_flatten_spec(tuple, _tuple_flatten_spec, _tuple_flatten_spec_exact_match)
+register_pytree_flatten_spec(namedtuple, _namedtuple_flatten_spec, _tuple_flatten_spec_exact_match)