[ez][BE] reuse formatting functions from utils in eval (#61)

Author: PaliC

BackendBench/eval.py

@@ -6,32 +6,20 @@
 
 
 from BackendBench.utils import uses_cuda_stream
+from BackendBench.utils import serialize_args
 
 logger = logging.getLogger(__name__)
 
 EXC_MSG = """
 Exception raised for {op}:
     args: {args}
-    kwargs: {kwargs}
     exc: {exc}
 """
 
 
-def format_tensor(t):
-    return f"{t.dtype}{list(t.shape)}"
-
-
-def format_args(args):
-    return [format_tensor(arg) if isinstance(arg, torch.Tensor) else arg for arg in args]
-
-
-def format_kwargs(kwargs):
-    return {k: format_tensor(v) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
-
-
 def format_exception(e, op, args, kwargs):
     op_name = getattr(op, "__name__", str(op))
-    return EXC_MSG.format(op=op_name, args=format_args(args), kwargs=format_kwargs(kwargs), exc=e)
+    return EXC_MSG.format(op=op_name, args=serialize_args(args, kwargs), exc=e)
 
 
 def allclose(a, b):
@@ -60,9 +48,7 @@ def eval_correctness_test(op, impl, test):
 def eval_correctness(op, impl, tests):
     correct, total = 0, 0
     for test in tests:
-        logging.debug(
-            f"Testing {op.__name__} with args {format_args(test.args)} and kwargs {format_kwargs(test.kwargs)}"
-        )
+        logging.debug(f"Testing {op.__name__} with args {serialize_args(test.args, test.kwargs)}")
         if eval_correctness_test(op, impl, test):
             correct += 1
         total += 1
@@ -88,7 +74,7 @@ def eval_performance(op, impl, tests):
     test_times = []
     for test in tests:
         logging.debug(
-            f"Benchmarking {op.__name__} with args {format_args(test.args)} and kwargs {format_kwargs(test.kwargs)}"
+            f"Benchmarking {op.__name__} with args {serialize_args(test.args, test.kwargs)}"
         )
         base_times.append(bench_fn(lambda: op(*test.args, **test.kwargs)))
         try:

test/test_eval.py

@@ -4,9 +4,6 @@
 try:
     import importlib.util
     from BackendBench.eval import (
-        format_tensor,
-        format_args,
-        format_kwargs,
         format_exception,
         allclose,
         eval_correctness_test,
@@ -24,37 +21,6 @@
 
 
 class TestFormatFunctions:
-    def test_format_tensor(self):
-        tensor = torch.randn(2, 3, 4, dtype=torch.float32)
-        formatted = format_tensor(tensor)
-        assert formatted == "torch.float32[2, 3, 4]"
-
-        tensor_int = torch.randint(0, 10, (5, 5), dtype=torch.int64)
-        formatted_int = format_tensor(tensor_int)
-        assert formatted_int == "torch.int64[5, 5]"
-
-    def test_format_args(self):
-        tensor1 = torch.randn(2, 3)
-        tensor2 = torch.randn(3, 4)
-        scalar = 2.5
-
-        args = [tensor1, scalar, tensor2]
-        formatted = format_args(args)
-
-        assert len(formatted) == 3
-        assert formatted[0] == "torch.float32[2, 3]"
-        assert formatted[1] == 2.5
-        assert formatted[2] == "torch.float32[3, 4]"
-
-    def test_format_kwargs(self):
-        tensor = torch.randn(2, 3)
-        kwargs = {"input": tensor, "dim": 1, "keepdim": True}
-
-        formatted = format_kwargs(kwargs)
-        assert formatted["input"] == "torch.float32[2, 3]"
-        assert formatted["dim"] == 1
-        assert formatted["keepdim"] is True
-
     def test_format_exception(self):
         op = torch.ops.aten.relu.default
         args = [torch.randn(2, 3)]