Refactor IR value creation in tests to use ir.Value for consistency and improved clarity

Author: whyvineet

onnxscript/rewriter/onnxruntime/bfloat16_utils/bfloat16_converter_test.py

@@ -6,20 +6,23 @@
 import onnx
 import onnx.checker
 import onnx.shape_inference
+import onnx_ir as ir
 import onnxruntime
 
-from onnxscript import ir
 from onnxscript.rewriter.onnxruntime.bfloat16_utils import bfloat16_converter
 
 
 class Bfloat16ConversionTest(unittest.TestCase):
     def setUp(self) -> None:
-        self.v0 = ir.val(name="v0", shape=ir.Shape([2, 3, 4]))
-        self.v0.dtype = ir.DataType.BFLOAT16
-        self.v1 = ir.val(name="v1", shape=ir.Shape([2, 3, 4]))
-        self.v1.dtype = ir.DataType.BFLOAT16
-        self.v2 = ir.val(name="v2", shape=ir.Shape([2, 3, 4]))
-        self.v2.dtype = ir.DataType.BFLOAT16
+        self.v0 = ir.Value(
+            name="v0", shape=ir.Shape([2, 3, 4]), type=ir.TensorType(ir.DataType.BFLOAT16)
+        )
+        self.v1 = ir.Value(
+            name="v1", shape=ir.Shape([2, 3, 4]), type=ir.TensorType(ir.DataType.BFLOAT16)
+        )
+        self.v2 = ir.Value(
+            name="v2", shape=ir.Shape([2, 3, 4]), type=ir.TensorType(ir.DataType.BFLOAT16)
+        )
 
         self.add_node = ir.Node("", "Add", inputs=(self.v0, self.v1), num_outputs=1)
         self.add_node.outputs[0].dtype = ir.DataType.BFLOAT16

onnxscript/rewriter/rules/common/_basic_rules_test.py

@@ -8,11 +8,11 @@
 import numpy as np
 import onnx
 import onnx.reference
+import onnx_ir as ir
 import parameterized
 
 import onnxscript
 import onnxscript.onnx_types as ot
-from onnxscript import ir
 from onnxscript.onnx_opset import opset18
 from onnxscript.rewriter import MatchingTracer, testing
 from onnxscript.rewriter import pattern as orp
@@ -421,14 +421,18 @@ def _convert_shape(shape, name):
             if isinstance(shape, np.ndarray):
                 shape = tape.initializer(ir.Tensor(shape, name=name))
             elif isinstance(shape, (list, tuple)):
-                shape = ir.val(name, ir.DataType.INT64, ir.Shape(shape))
+                shape = ir.Value(
+                    name=name, type=ir.TensorType(ir.DataType.INT64), shape=ir.Shape(shape)
+                )
                 tape.graph_like.inputs.append(shape)
             else:
                 raise TypeError(f"Unsupported type {type(shape)} for shape.")
             return shape
 
-        x = ir.val("X", ir.DataType.FLOAT, ir.Shape(input_shape))
-        y = ir.val("Y", ir.DataType.FLOAT)
+        x = ir.Value(
+            name="X", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(input_shape)
+        )
+        y = ir.Value(name="Y", type=ir.TensorType(ir.DataType.FLOAT))
         tape = ir.tape.Tape(ir.Graph([x], [y], nodes=[], opset_imports={"": 20}))
 
         # Build the graph.
@@ -554,8 +558,10 @@ def test_unsupported_reshape_reshape(self, shape2, error_msg):
 class Flatten2ReshapeTest(unittest.TestCase):
     @staticmethod
     def create_model(input_shape, axis=1):
-        x = ir.val("X", ir.DataType.FLOAT, ir.Shape(input_shape))
-        y = ir.val("Y", ir.DataType.FLOAT)
+        x = ir.Value(
+            name="X", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(input_shape)
+        )
+        y = ir.Value(name="Y", type=ir.TensorType(ir.DataType.FLOAT))
         tape = ir.tape.Tape(ir.Graph([x], [y], nodes=[], opset_imports={"": 20}))
 
         # Build the graph.

onnxscript/rewriter/rules/common/_fuse_pad_into_conv_test.py

@@ -61,13 +61,13 @@ def build_model(
 
         # Register operations in the tape
         idtype = ir.DataType.UINT8 if op_type == "ConvInteger" else ir.DataType.FLOAT
-        x = ir.val("X", shape=input_shape, type=ir.TensorType(idtype))
+        x = ir.Value(name="X", shape=input_shape, type=ir.TensorType(idtype))
         y = tape.op("Pad", inputs=[x, *pad_inputs], attributes=pad_attributes)
         y = tape.op(
             op_type,
             inputs=[y, self.get_conv_weights(weight_shape, tape)],
             attributes=conv_attributes,
-            output=ir.val("Y", shape=output_shape, type=ir.TensorType(x.dtype)),
+            output=ir.Value(name="Y", shape=output_shape, type=ir.TensorType(x.dtype)),
         )
         if op_type == "ConvInteger":
             y.dtype = ir.DataType.INT32
@@ -290,12 +290,12 @@ def build_model(
                 raise ValueError(f"Unsupported type for pad input ({x}): {type(x)}.")
 
         # Register operations in the tape
-        x = ir.val("X", shape=input_shape, type=ir.TensorType(ir.DataType.FLOAT))
+        x = ir.Value(name="X", shape=input_shape, type=ir.TensorType(ir.DataType.FLOAT))
         y = tape.op(
             "Conv",
             inputs=[x, *conv_inputs],
             attributes=conv_attributes,
-            output=ir.val("Y", shape=output_shape, type=x.type),
+            output=ir.Value(name="Y", shape=output_shape, type=x.type),
         )
 
         # Build the model

onnxscript/rewriter/rules/common/_matmul_add_to_gemm_test.py

@@ -5,10 +5,10 @@
 
 import numpy as np
 import onnx
+import onnx_ir as ir
 from onnx_ir.passes.common import onnx_checker, shape_inference
 from parameterized import parameterized
 
-from onnxscript import ir
 from onnxscript.rewriter import MatchingTracer, MatchStatus, testing
 from onnxscript.rewriter.rules.common import _matmul_add_to_gemm
 
@@ -46,10 +46,10 @@ def get_test_model(
         bias_shape = weight_shape[0] if transB else weight_shape[-1]
         output_shape = ir.Shape(("?",) * input_shape.rank())
 
-        x = ir.val("X", shape=input_shape, type=ir.TensorType(ir.DataType.FLOAT))
+        x = ir.Value(name="X", shape=input_shape, type=ir.TensorType(ir.DataType.FLOAT))
 
         if weight_as_inputs:
-            w = ir.val("W", shape=weight_shape, type=ir.TensorType(ir.DataType.FLOAT))
+            w = ir.Value(name="W", shape=weight_shape, type=ir.TensorType(ir.DataType.FLOAT))
             inputs.append(w)
         else:
             w = ir.tensor(
@@ -58,8 +58,8 @@ def get_test_model(
             w = tape.initializer(w)
 
         if bias_as_inputs:
-            b = ir.val(
-                "B", shape=ir.Shape([bias_shape]), type=ir.TensorType(ir.DataType.FLOAT)
+            b = ir.Value(
+                name="B", shape=ir.Shape([bias_shape]), type=ir.TensorType(ir.DataType.FLOAT)
             )
             inputs.append(b)
         else:
@@ -77,7 +77,9 @@ def get_test_model(
         y = tape.op(
             "Add",
             inputs=[y, b],
-            output=ir.val("Y", shape=output_shape, type=ir.TensorType(ir.DataType.FLOAT)),
+            output=ir.Value(
+                name="Y", shape=output_shape, type=ir.TensorType(ir.DataType.FLOAT)
+            ),
         )
 
         # Build the model

onnxscript/rewriter/rules/common/_remove_zero_bias_test.py

@@ -61,9 +61,17 @@ def _apply_rule_and_check_optimization(
             assert attr_name in target_node.attributes, f"Attribute {attr_name} was lost"
             original_value = attr_value.value
             new_value = target_node.attributes[attr_name].value
-            assert new_value == original_value, (
-                f"Attribute {attr_name} value changed from {original_value} to {new_value}"
-            )
+            # Convert both to same type for comparison to handle list vs tuple differences
+            if isinstance(original_value, (list, tuple)) and isinstance(
+                new_value, (list, tuple)
+            ):
+                assert list(original_value) == list(new_value), (
+                    f"Attribute {attr_name} value changed from {original_value} to {new_value}"
+                )
+            else:
+                assert new_value == original_value, (
+                    f"Attribute {attr_name} value changed from {original_value} to {new_value}"
+                )
 
     # Check specific expected attributes if provided
     if expected_attributes:
@@ -73,9 +81,17 @@ def _apply_rule_and_check_optimization(
             )
             actual_attr = target_node.attributes[attr_name]
             actual_value = actual_attr.value
-            assert actual_value == expected_value, (
-                f"Expected attribute {attr_name} to be {expected_value}, got {actual_value}"
-            )
+            # Convert both to same type for comparison to handle list vs tuple differences
+            if isinstance(actual_value, (list, tuple)) and isinstance(
+                expected_value, (list, tuple)
+            ):
+                assert list(actual_value) == list(expected_value), (
+                    f"Expected attribute {attr_name} to be {expected_value}, got {actual_value}"
+                )
+            else:
+                assert actual_value == expected_value, (
+                    f"Expected attribute {attr_name} to be {expected_value}, got {actual_value}"
+                )
 
     # Compare outputs to ensure correctness (only for supported input types)
     if expected_count > 0: