Fix rebase bug with ttrt identity; add atol/rtol (#363)

tenstorrent · Aug 12, 2024 · 6358122 · 6358122
1 parent 0b393ec
commit 6358122
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 9 deletions.
diff --git a/runtime/tools/python/ttrt/__init__.py b/runtime/tools/python/ttrt/__init__.py
@@ -91,6 +91,18 @@ def main():
         action="store_true",
         help="Do a golden identity test on the output tensors",
     )
+    run_parser.add_argument(
+        "--rtol",
+        default=1e-05,
+        type=float,
+        help="rtol for golden test",
+    )
+    run_parser.add_argument(
+        "--atol",
+        default=1e-08,
+        type=float,
+        help="atol for golden test",
+    )
     run_parser.add_argument(
         "--seed",
         default=0,

diff --git a/runtime/tools/python/ttrt/common/api.py b/runtime/tools/python/ttrt/common/api.py
@@ -45,6 +45,21 @@ def volume(shape):
 }
 init_fns = sorted(list(init_fns_map.keys()))
 
+
+def check_identity(args, binary_name, torch_inputs, torch_outputs):
+    import torch
+
+    for i, o in zip(torch_inputs, torch_outputs):
+        if not torch.allclose(i, o, rtol=args.rtol, atol=args.atol):
+            print(
+                "Failed: inputs and outputs do not match in binary",
+                binary_name,
+            )
+            print(torch.abs(i - o))
+        else:
+            print("Passed:", binary_name)
+
+
 #######################################################################################
 ########################################**API**########################################
 #######################################################################################
@@ -252,15 +267,12 @@ def run(args):
             ttrt.runtime.wait(event)
             print("outputs:\n", torch_outputs)
             if args.identity:
-                for i, o in zip(torch_inputs[binary_name], torch_outputs[binary_name]):
-                    if not torch.allclose(i, o):
-                        print(
-                            "Failed: inputs and outputs do not match in binary",
-                            binary_name,
-                        )
-                        print(i - o)
-                    else:
-                        print("Passed:", binary_name)
+                check_identity(
+                    args,
+                    binary_name,
+                    torch_inputs[binary_name][program_index],
+                    torch_outputs[binary_name][program_index],
+                )
 
     # save artifacts
     if arg_save_artifacts: