Make metadata serialization more strict (#124411)

Summary: When I was debugging an issue, this silent error makes the debugging harder. It is better to error earlier with more descriptive error message.

Test Plan: None

Differential Revision: D56312433

Pull Request resolved: #124411
Approved by: https://github.com/zhxchen17

torch/export/_trace.py

@@ -589,10 +589,13 @@ def make_argument_spec(i, node) -> ArgumentSpec:
             return CustomObjArgument(
                 name=node.name, class_fqn=val._type().qualified_name()  # type: ignore[attr-defined]
             )
-        else:
-            # TODO: this branch is likely wrong, all permissible ConstantArgument type
-            # should have been handled already
+        elif isinstance(val, (int, bool, str, float, type(None))):
             return ConstantArgument(name=node.name, value=val)
+        else:
+            raise AssertionError(
+                f"Encountered an unsupported object of type {type(val)} "
+                f"while writing the metadata for exported program"
+            )
 
     input_specs, output_specs = _sig_to_specs(
         user_inputs=set(graph_signature.user_inputs),

torch/export/graph_signature.py

@@ -41,7 +41,7 @@ class CustomObjArgument:
 @dataclasses.dataclass
 class ConstantArgument:
     name: str
-    value: Union[int, float, bool, None]
+    value: Union[int, float, bool, str, None]
 
 
 ArgumentSpec = Union[