address reviews

Author: titaiwangms

onnxscript/ir/passes/common/constant_manipulation.py

@@ -87,33 +87,28 @@ def _constant_node_attribute_to_tensor(
             )
             return None
 
-        # Dispatch table for attribute-to-tensor conversion
-        tensor_converters = {
-            "value": lambda: attr_value.as_tensor(),  # pylint: disable=unnecessary-lambda`
-            "value_int": lambda: ir.tensor(
+        if attr_name == "value":
+            tensor = attr_value.as_tensor()  # type: ignore[union-attr]
+        elif attr_name == "value_int":
+            tensor = ir.tensor(
                 attr_value.as_int(), dtype=ir.DataType.INT64, name=initializer_name
-            ),
-            "value_ints": lambda: ir.tensor(
+            )
+        elif attr_name == "value_ints":
+            tensor = ir.tensor(
                 attr_value.as_ints(), dtype=ir.DataType.INT64, name=initializer_name
-            ),
-            "value_float": lambda: ir.tensor(
+            )
+        elif attr_name == "value_float":
+            tensor = ir.tensor(
                 attr_value.as_float(), dtype=ir.DataType.FLOAT, name=initializer_name
-            ),
-            "value_floats": lambda: ir.tensor(
+            )
+        elif attr_name == "value_floats":
+            tensor = ir.tensor(
                 attr_value.as_floats(), dtype=ir.DataType.FLOAT, name=initializer_name
-            ),
-            "value_string": lambda: ir.StringTensor(
-                np.array(attr_value.value, dtype=np.bytes_), name=initializer_name
-            ),
-            "value_strings": lambda: ir.StringTensor(
+            )
+        elif attr_name in ("value_string", "value_strings"):
+            tensor = ir.StringTensor(
                 np.array(attr_value.value, dtype=np.bytes_), name=initializer_name
-            ),
-        }
-        converter = tensor_converters.get(attr_name)
-        if converter is None:
-            logger.debug(
-                "Unsupported constant node attribute '%s' in node '%s'", attr_name, node.name
             )
-            return None
-
-        return converter()  # type: ignore[return-value]
+        else:
+            tensor = None
+        return tensor  # type: ignore[return-value]

onnxscript/ir/passes/common/constant_manipulation_test.py

@@ -21,7 +21,7 @@ class TestLiftConstantsToInitializersPass(unittest.TestCase):
         ]
     )
     def test_pass_with_lifting_float_and_int_constants_to_initializers(
-        self, ir_dtype, lift_all_constants
+        self, ir_dtype: ir.DataType, lift_all_constants: bool
     ):
         inputs = [
             ir.Value(name="input_a", type=ir.TensorType(ir_dtype), shape=ir.Shape((2, 3))),
@@ -80,7 +80,7 @@ def test_pass_with_lifting_float_and_int_constants_to_initializers(
         ]
     )
     def test_pass_with_lifting_constants_to_initializers_within_subgraph(
-        self, lift_all_constants
+        self, lift_all_constants: bool
     ):
         input_value = ir.Value(
             name="input", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3))
@@ -167,7 +167,11 @@ def test_pass_with_lifting_constants_to_initializers_within_subgraph(
         ]
     )
     def test_pass_with_lifting_constants_to_initializers_with_floats_ints_strings(
-        self, value, constant_attribute, np_dtype, lift_all_constants
+        self,
+        value: float | int | str | list[float] | list[int] | list[str],
+        constant_attribute: str,
+        np_dtype: type[np.dtype],
+        lift_all_constants: bool,
     ):
         input_value = ir.Value(
             name="input", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3))