From b596493113fadda9c87be9b521006e19eaf5337d Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 9 May 2023 09:57:47 -0500 Subject: [PATCH] [Testing] Use TVMScript's "name" argument for error messages Prior to this commit, the `tvm.testing.CompareBeforeAfter` utility performed string replacement to define the names of the before/after/expected examples. Since https://github.com/apache/tvm/pull/13934 allowed the name to be explicitly passed to the printer, this should be used instead. --- python/tvm/testing/utils.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py index 399adf27e4e5..884a885fb21d 100644 --- a/python/tvm/testing/utils.py +++ b/python/tvm/testing/utils.py @@ -2065,13 +2065,6 @@ def _is_method(func): def test_compare(self, before, expected, transform): """Unit test to compare the expected TIR PrimFunc to actual""" - def pprint(name, obj): - script = obj.script() - if isinstance(obj, tvm.IRModule): - return script.replace("class Module", f"class {name}") - else: - return script.replace("def func", f"def {name}") - if inspect.isclass(expected) and issubclass(expected, Exception): with pytest.raises(expected): after = transform(before) @@ -2079,8 +2072,8 @@ def pprint(name, obj): # This portion through pytest.fail isn't strictly # necessary, but gives a better error message that # includes the before/after. - before_str = pprint("before", before) - after_str = pprint("after", after) + before_str = before.script(name="before") + after_str = after.script(name="after") pytest.fail( msg=( @@ -2095,9 +2088,9 @@ def pprint(name, obj): try: tvm.ir.assert_structural_equal(after, expected) except ValueError as err: - before_str = pprint("before", before) - after_str = pprint("after", after) - expected_str = pprint("expected", expected) + before_str = before.script(name="before") + after_str = after.script(name="after") + expected_str = expected.script(name="expected") raise ValueError( f"TIR after transformation did not match expected:\n" f"{before_str}\n{after_str}\n{expected_str}"