Added better error messages for type validators (stack trace)

Author: cptspacemanspiff

exir/verification/arg_validator.py

@@ -37,7 +37,7 @@ class EdgeOpArgValidator(torch.fx.Interpreter):
 
     def __init__(self, graph_module: torch.fx.GraphModule) -> None:
         super().__init__(graph_module)
-        self.violating_ops: Dict[EdgeOpOverload, Dict[str, Optional[torch.dtype]]] = (
+        self.violating_ops: Dict[EdgeOpOverload, Tuple[Dict[str, Optional[torch.dtype]], torch.fx.Node]] = (
             defaultdict(dict)
         )
 
@@ -125,5 +125,5 @@ def call_function(  # noqa: C901  # pyre-fixme[14]
 
         valid = target._schema.dtype_constraint.validate(tensor_arg_types)
         if not valid:
-            self.violating_ops[target] = tensor_arg_types
+            self.violating_ops[target] = (tensor_arg_types, self.node)
         return super().call_function(target, args, kwargs)  # pyre-fixme[6]

exir/verification/verifier.py

@@ -189,9 +189,15 @@ def _check_tensor_args_matching_op_allowed_dtype(gm: GraphModule) -> None:
         return
 
     if validator.violating_ops:
+        error_msg = ""
+        for op, node in validator.violating_ops.items():
+            # error_msg += f"#####################################################\n"
+            error_msg += f"Operator: {op} with args: {node[0]}\n"
+            error_msg += f"stack trace: {node[1].stack_trace}\n\n"
+            # error_msg += f"#####################################################\n"
         raise SpecViolationError(
-            f"These operators are taking Tensor inputs with mismatched dtypes: {validator.violating_ops}"
-            "Please make sure the dtypes of the Tensor inputs are the same as the dtypes of the corresponding "
+            f"These operators are taking Tensor inputs with mismatched dtypes:\n{error_msg}\n"
+            "Please make sure the dtypes of the Tensor inputs are the same as the dtypes of the corresponding outputs."
         )