Add testing

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

src/onnx_ir/testing.py

@@ -1,3 +1,196 @@
 # Copyright (c) ONNX Project Contributors
 # SPDX-License-Identifier: Apache-2.0
 """Utilities for testing."""
+
+from __future__ import annotations
+
+__all__ = [
+    "assert_onnx_proto_equal",
+]
+
+import difflib
+import math
+from typing import Any, Collection, Sequence
+
+import google.protobuf.message
+import onnx
+
+
+def _opset_import_key(opset_import: onnx.OperatorSetIdProto) -> tuple[str, int]:
+    return (opset_import.domain, opset_import.version)
+
+
+def _value_info_key(value_info: onnx.ValueInfoProto) -> str:
+    return value_info.name
+
+
+def _function_key(function: onnx.FunctionProto) -> tuple[str, str, str]:
+    return (function.domain, function.name, getattr(function, "overload", ""))
+
+
+def _find_duplicates(with_duplicates: Collection[Any]) -> list[Any]:
+    """Return a list of duplicated elements in a collection."""
+    seen = set()
+    duplicates = []
+    for x in with_duplicates:
+        if x in seen:
+            duplicates.append(x)
+        seen.add(x)
+    return duplicates
+
+
+def assert_onnx_proto_equal(
+    actual: google.protobuf.message.Message | Any,
+    expected: google.protobuf.message.Message | Any,
+    ignore_initializer_value_proto: bool = False,
+) -> None:
+    """Assert that two ONNX protos are equal.
+
+    Equality is defined as having the same fields with the same values. When
+    a field takes the default value, it is considered equal to the field
+    not being set.
+
+    Sequential fields with name `opset_import`, `value_info`, and `functions` are
+    compared disregarding the order of their elements.
+
+    Args:
+        actual: The first ONNX proto.
+        expected: The second ONNX proto.
+        ignore_initializer_value_proto: Ignore value protos for initializers if there
+            are extra ones in the actual proto.
+    """
+    assert type(actual) is type(expected), (
+        f"Type not equal: {type(actual)} != {type(expected)}"
+    )
+
+    a_fields = {field.name: value for field, value in actual.ListFields()}
+    b_fields = {field.name: value for field, value in expected.ListFields()}
+    all_fields = sorted(set(a_fields.keys()) | set(b_fields.keys()))
+    if isinstance(actual, onnx.GraphProto) and isinstance(expected, onnx.GraphProto):
+        actual_initializer_names = {i.name for i in actual.initializer}
+        expected_initializer_names = {i.name for i in expected.initializer}
+    else:
+        actual_initializer_names = set()
+        expected_initializer_names = set()
+
+    # Record and report all errors
+    errors = []
+    for field in all_fields:  # pylint: disable=too-many-nested-blocks
+        # Obtain the default value if the field is not set. This way we can compare the two fields.
+        a_value = getattr(actual, field)
+        b_value = getattr(expected, field)
+        if (
+            isinstance(a_value, Sequence)
+            and isinstance(b_value, Sequence)
+            and not isinstance(a_value, (str, bytes))
+            and not isinstance(b_value, (str, bytes))
+        ):
+            # Check length first
+            a_keys: list[Any] = []
+            b_keys: list[Any] = []
+            if field == "opset_import":
+                a_value = sorted(a_value, key=_opset_import_key)
+                b_value = sorted(b_value, key=_opset_import_key)
+                a_keys = [_opset_import_key(opset_import) for opset_import in a_value]
+                b_keys = [_opset_import_key(opset_import) for opset_import in b_value]
+            elif field == "value_info":
+                if (
+                    ignore_initializer_value_proto
+                    and isinstance(actual, onnx.GraphProto)
+                    and isinstance(expected, onnx.GraphProto)
+                ):
+                    # Filter out initializers from the value_info list
+                    a_value = [
+                        value_info
+                        for value_info in a_value
+                        if value_info.name not in actual_initializer_names
+                    ]
+                    b_value = [
+                        value_info
+                        for value_info in b_value
+                        if value_info.name not in expected_initializer_names
+                    ]
+                a_value = sorted(a_value, key=_value_info_key)
+                b_value = sorted(b_value, key=_value_info_key)
+                a_keys = [_value_info_key(value_info) for value_info in a_value]
+                b_keys = [_value_info_key(value_info) for value_info in b_value]
+            elif field == "functions":
+                a_value = sorted(a_value, key=_function_key)
+                b_value = sorted(b_value, key=_function_key)
+                a_keys = [_function_key(functions) for functions in a_value]
+                b_keys = [_function_key(functions) for functions in b_value]
+
+            if a_keys != b_keys:
+                keys_only_in_actual = set(a_keys) - set(b_keys)
+                keys_only_in_expected = set(b_keys) - set(a_keys)
+                error_message = (
+                    f"Field {field} not equal: keys_only_in_actual={keys_only_in_actual}, keys_only_in_expected={keys_only_in_expected}. "
+                    f"Field type: {type(a_value)}. "
+                    f"Duplicated a_keys: {_find_duplicates(a_keys)}, duplicated b_keys: {_find_duplicates(b_keys)}"
+                )
+                errors.append(error_message)
+            elif len(a_value) != len(b_value):
+                error_message = (
+                    f"Field {field} not equal: len(a)={len(a_value)}, len(b)={len(b_value)} "
+                    f"Field type: {type(a_value)}"
+                )
+                errors.append(error_message)
+            else:
+                # Check every element
+                for i in range(len(a_value)):  # pylint: disable=consider-using-enumerate
+                    actual_value_i = a_value[i]
+                    expected_value_i = b_value[i]
+                    if isinstance(
+                        actual_value_i, google.protobuf.message.Message
+                    ) and isinstance(expected_value_i, google.protobuf.message.Message):
+                        try:
+                            assert_onnx_proto_equal(
+                                actual_value_i,
+                                expected_value_i,
+                                ignore_initializer_value_proto=ignore_initializer_value_proto,
+                            )
+                        except AssertionError as e:
+                            error_message = f"Field {field} index {i} in sequence not equal. type(actual_value_i): {type(actual_value_i)}, type(expected_value_i): {type(expected_value_i)}, actual_value_i: {actual_value_i}, expected_value_i: {expected_value_i}"
+                            error_message = (
+                                str(e) + "\n\nCaused by the above error\n\n" + error_message
+                            )
+                            errors.append(error_message)
+                    elif actual_value_i != expected_value_i:
+                        if (
+                            isinstance(actual_value_i, float)
+                            and isinstance(expected_value_i, float)
+                            and math.isnan(actual_value_i)
+                            and math.isnan(expected_value_i)
+                        ):
+                            # Consider NaNs equal
+                            continue
+                        error_message = f"Field {field} index {i} in sequence not equal. type(actual_value_i): {type(actual_value_i)}, type(expected_value_i): {type(expected_value_i)}"
+                        for line in difflib.ndiff(
+                            str(actual_value_i).splitlines(),
+                            str(expected_value_i).splitlines(),
+                        ):
+                            error_message += "\n" + line
+                        errors.append(error_message)
+        elif isinstance(a_value, google.protobuf.message.Message) and isinstance(
+            b_value, google.protobuf.message.Message
+        ):
+            assert_onnx_proto_equal(
+                a_value, b_value, ignore_initializer_value_proto=ignore_initializer_value_proto
+            )
+        elif a_value != b_value:
+            if (
+                isinstance(a_value, float)
+                and isinstance(b_value, float)
+                and math.isnan(a_value)
+                and math.isnan(b_value)
+            ):
+                # Consider NaNs equal
+                continue
+            error_message = (
+                f"Field {field} not equal. field_actual: {a_value}, field_expected: {b_value}"
+            )
+            errors.append(error_message)
+    if errors:
+        raise AssertionError(
+            f"Protos not equal: {type(actual)} != {type(expected)}\n" + "\n".join(errors)
+        )