onnx 1.17.0__cp312-cp312-win32.whl → 1.19.0__cp312-cp312-win32.whl

onnx/test/shape_inference_test.py

@@ -4,9 +4,10 @@
 
 from __future__ import annotations
 
+import contextlib
 import itertools
 import unittest
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pytest
@@ -37,6 +38,7 @@ from onnx.defs import (
 )
 from onnx.helper import (
     make_empty_tensor_value_info,
+    make_graph,
     make_node,
     make_opsetid,
     make_tensor,
@@ -45,6 +47,9 @@ from onnx.helper import (
 )
 from onnx.parser import parse_graph
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
 
 def get_available_versions(schema: OpSchema) -> set[int]:
     versions: set[int] = set()
@@ -161,8 +166,8 @@ class TestShapeInferenceHelper(unittest.TestCase):
         vis = [x for x in graph.value_info if x.name not in names_in_vis] + vis
         inferred_model = self._inferred(graph_or_model, **kwargs)
         inferred_vis = list(inferred_model.graph.value_info)
-        vis = sorted(vis, key=lambda x: x.name)  # type: ignore[no-any-return]
-        inferred_vis = sorted(inferred_vis, key=lambda x: x.name)  # type: ignore
+        vis = sorted(vis, key=lambda x: x.name)
+        inferred_vis = sorted(inferred_vis, key=lambda x: x.name)
         assert len(vis) == len(inferred_vis)
         for v, inferred_v in zip(vis, inferred_vis):
             self._compare_value_infos(v.type, inferred_v.type)
@@ -188,13 +193,13 @@ class TestShapeInferenceHelper(unittest.TestCase):
                     inferred_dim = inferred_vi_type.tensor_type.shape.dim[dim_i]
                     # if it is a symbolic shape, make sure the inferred symbol has generated (dim_param)
                     if dim.dim_param:
-                        assert (
-                            dim.dim_param == inferred_dim.dim_param
-                        ), f"\n{vi_type}\n{inferred_vi_type}\n"
+                        assert dim.dim_param == inferred_dim.dim_param, (
+                            f"\n{vi_type}\n{inferred_vi_type}\n"
+                        )
                     else:
-                        assert (
-                            dim.dim_value == inferred_dim.dim_value
-                        ), f"\n{vi_type}\n{inferred_vi_type}\n"
+                        assert dim.dim_value == inferred_dim.dim_value, (
+                            f"\n{vi_type}\n{inferred_vi_type}\n"
+                        )
         elif vi_type.HasField("sequence_type"):
             assert inferred_vi_type.HasField("sequence_type")
             vi = vi_type.sequence_type.elem_type
@@ -284,7 +289,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             [("X", TensorProto.FLOAT, (2, 3, 4))],
             [make_node("Transpose", ["X"], ["Y"], perm=[1, 0, 2])],
             [make_tensor_value_info("Y", TensorProto.UNDEFINED, (3, "a", "b"))],
-        )  # type: ignore
+        )
         self._assert_inferred(
             graph,
             [make_tensor_value_info("Y", TensorProto.FLOAT, (3, 2, 4))],
@@ -976,6 +981,22 @@ class TestShapeInference(TestShapeInferenceHelper):
                 opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
             )
 
+    @parameterized.expand(all_versions_for("RMSNormalization"))
+    def test_rms_normalization(self, _, version) -> None:
+        graph = self._make_graph(
+            [
+                ("X", TensorProto.FLOAT, ("N", "C", "H", "W")),
+                ("scale", TensorProto.FLOAT, ("H", "W")),
+            ],
+            [make_node("RMSNormalization", ["X", "scale"], ["y"], axis=2)],
+            [],
+        )
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info("y", TensorProto.FLOAT, ("N", "C", "H", "W"))],
+            opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
+        )
+
     @parameterized.expand(all_versions_for("Resize"))
     def test_resize_size_axes_2_3(self, _, version) -> None:
         self.skipIf(version < 18, "axes is from Version 18")
@@ -1389,7 +1410,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph,
-            [make_tensor_value_info("y", TensorProto.FLOAT, (2, 3))],  # type: ignore
+            [make_tensor_value_info("y", TensorProto.FLOAT, (2, 3))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
         )
 
@@ -1402,7 +1423,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph,
-            [make_tensor_value_info("y", TensorProto.FLOAT, (4, 1, 2, 5))],  # type: ignore
+            [make_tensor_value_info("y", TensorProto.FLOAT, (4, 1, 2, 5))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
         )
 
@@ -1428,7 +1449,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph,
-            [make_tensor_value_info("y", TensorProto.FLOAT, (2, 2))],  # type: ignore
+            [make_tensor_value_info("y", TensorProto.FLOAT, (2, 2))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
         )
 
@@ -1441,7 +1462,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph,
-            [make_tensor_value_info("y", TensorProto.FLOAT, (2, 3))],  # type: ignore
+            [make_tensor_value_info("y", TensorProto.FLOAT, (2, 3))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
         )
 
@@ -1470,7 +1491,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             graph,
             [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
-        )  # type: ignore
+        )
 
     @parameterized.expand(all_versions_for("Scatter"))
     def test_scatter_axis1(self, _, version) -> None:
@@ -1497,7 +1518,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             graph,
             [make_tensor_value_info("y", TensorProto.FLOAT, (1, 5))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
-        )  # type: ignore
+        )
 
     @parameterized.expand(all_versions_for("ScatterElements"))
     def test_scatter_elements(self, _, version) -> None:
@@ -1512,7 +1533,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph,
-            [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))],  # type: ignore
+            [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
         )
 
@@ -1529,7 +1550,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph,
-            [make_tensor_value_info("y", TensorProto.FLOAT, (1, 5))],  # type: ignore
+            [make_tensor_value_info("y", TensorProto.FLOAT, (1, 5))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
         )
 
@@ -1546,7 +1567,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph,
-            [make_tensor_value_info("y", TensorProto.FLOAT, (4, 5, 6))],  # type: ignore
+            [make_tensor_value_info("y", TensorProto.FLOAT, (4, 5, 6))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
         )
 
@@ -1574,7 +1595,34 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("y", TensorProto.FLOAT, None),
             ],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
-        )  # type: ignore
+        )
+
+    def test_tensor_scatter(self) -> None:
+        graph = self._make_graph(
+            [
+                ("past_cache", TensorProto.FLOAT, (2, 8, 128, 64)),
+                ("update", TensorProto.FLOAT, (2, 8, 10, 64)),
+                ("write_indices", TensorProto.INT64, (2,)),
+            ],
+            [
+                make_node(
+                    "TensorScatter",
+                    ["past_cache", "update", "write_indices"],
+                    ["present_cache"],
+                    axis=2,
+                )
+            ],
+            [],
+        )
+        self._assert_inferred(
+            graph,
+            [
+                make_tensor_value_info(
+                    "present_cache", TensorProto.FLOAT, (2, 8, 128, 64)
+                )
+            ],
+            opset_imports=[helper.make_opsetid(ONNX_DOMAIN, 24)],
+        )
 
     @parameterized.expand(all_versions_for("Squeeze"))
     def test_squeeze(self, _, version) -> None:
@@ -1680,6 +1728,23 @@ class TestShapeInference(TestShapeInferenceHelper):
             graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 2))]
         )
 
+    def test_squeeze_no_axes_dynamic_input_opset11(self) -> None:
+        graph = self._make_graph(
+            [
+                ("x", TensorProto.FLOAT, (1, 3, 1, None, 2, 1)),
+            ],
+            [make_node("Squeeze", ["x"], "y")],
+            [],
+        )
+        operatorsetid = OperatorSetIdProto()
+        operatorsetid.domain = ""
+        operatorsetid.version = 11
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info("y", TensorProto.FLOAT, None)],
+            opset_imports=[operatorsetid],
+        )
+
     def test_unsqueeze_regular(self) -> None:
         graph = self._make_graph(
             [("x", TensorProto.FLOAT, (3, 2)), ("axes", TensorProto.INT64, (4,))],
@@ -1776,7 +1841,9 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor("ends", TensorProto.INT64, (3,), (1, 1, 1)),
             ],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (1, None, 1))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (1, None, 1))]
+        )
 
     def test_slice_with_input_shape_steps(self) -> None:
         graph = self._make_graph(
@@ -1900,7 +1967,9 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor("axes", TensorProto.INT64, (2,), (0, 1)),
             ],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (2, 1))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (2, 1))]
+        )
 
     def test_slice_negative_start(self) -> None:
         graph = self._make_graph(
@@ -1920,7 +1989,9 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor("axes", TensorProto.INT64, (2,), (0, 1)),
             ],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (2, 2))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (2, 2))]
+        )
 
     def test_slice_negative_step(self) -> None:
         graph = self._make_graph(
@@ -1942,7 +2013,9 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor("steps", TensorProto.INT64, (2,), (1, -1)),
             ],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (2, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (2, 3))]
+        )
 
     def test_slice_variable_copy(self) -> None:
         graph = self._make_graph(
@@ -1960,7 +2033,9 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor("axes", TensorProto.INT64, (1,), (1,)),
             ],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, ("a", 1))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, ("a", 1))]
+        )
 
     def test_slice_variable_input_types(self) -> None:
         graph = self._make_graph(
@@ -2139,7 +2214,10 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Conv", ["x", "y"], "z", pads=[1, 1, 2, 0, 1, 2])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (30, 50, None, 6, 6))])  # type: ignore
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info("z", TensorProto.FLOAT, (30, 50, None, 6, 6))],
+        )
 
     def test_conv_partial_missing_weight_shape(self) -> None:
         graph = self._make_graph(
@@ -2154,6 +2232,45 @@ class TestShapeInference(TestShapeInferenceHelper):
             graph, [make_tensor_value_info("z", TensorProto.FLOAT, None)]
         )
 
+    def test_attention_4d(self) -> None:
+        graph = self._make_graph(
+            [
+                (
+                    "Q",
+                    TensorProto.FLOAT,
+                    ("B", "q_num_heads", "q_seq_length", "head_size"),
+                ),
+                (
+                    "K",
+                    TensorProto.FLOAT,
+                    ("B", "kv_num_heads", "kv_seq_len", "head_size"),
+                ),
+                (
+                    "V",
+                    TensorProto.FLOAT,
+                    ("B", "kv_num_heads", "kv_seq_len", "v_head_size"),
+                ),
+            ],
+            [
+                make_node(
+                    "Attention",
+                    ["Q", "K", "V"],
+                    ["Y"],
+                )
+            ],
+            [],
+        )
+        self._assert_inferred(
+            graph,
+            [
+                make_tensor_value_info(
+                    "Y",
+                    TensorProto.FLOAT,
+                    ("B", "q_num_heads", "q_seq_length", "v_head_size"),
+                )
+            ],
+        )
+
     def test_average_pool_auto_pads(self) -> None:
         graph = self._make_graph(
             [("x", TensorProto.FLOAT, (30, 4, 7, 6, 4))],
@@ -2231,12 +2348,14 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_sequence_value_info("in_sequence", TensorProto.FLOAT, (2, None, 4)),  # type: ignore
+                make_tensor_sequence_value_info(
+                    "in_sequence", TensorProto.FLOAT, (2, None, 4)
+                ),
                 make_tensor_sequence_value_info(
                     "output_sequence", TensorProto.FLOAT, (2, None, 4)
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_identity_optional(self) -> None:
         graph = self._make_graph(
@@ -2252,10 +2371,10 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                helper.make_value_info("in_optional", optional_type_proto),  # type: ignore
+                helper.make_value_info("in_optional", optional_type_proto),
                 helper.make_value_info("output_optional", optional_type_proto),
             ],
-        )  # type: ignore
+        )
 
     def test_identity_optional_sequence(self) -> None:
         graph = self._make_graph(
@@ -2281,11 +2400,11 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                helper.make_value_info("in_sequence", sequence_type_proto),  # type: ignore
-                helper.make_value_info("in_optional", optional_type_proto),  # type: ignore
+                helper.make_value_info("in_sequence", sequence_type_proto),
+                helper.make_value_info("in_optional", optional_type_proto),
                 helper.make_value_info("output_optional", optional_type_proto),
             ],
-        )  # type: ignore
+        )
 
     def test_add(self) -> None:
         graph = self._make_graph(
@@ -2441,7 +2560,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Bernoulli", ["x"], ["out"])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("out", TensorProto.FLOAT, (3, 4))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("out", TensorProto.FLOAT, (3, 4))]
+        )
 
     def test_bernoulli_with_dtype(self) -> None:
         graph = self._make_graph(
@@ -2456,7 +2577,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("out", TensorProto.DOUBLE, (2, 3, 4))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("out", TensorProto.DOUBLE, (2, 3, 4))]
+        )
 
     def _logical_binary_op(self, op: str, input_type: TensorProto.DataType) -> None:
         graph = self._make_graph(
@@ -2560,7 +2683,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Flatten", ["x"], ["z"], axis=2)],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (None, 20))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (None, 20))]
+        )
 
     def test_space_to_depth(self) -> None:
         b = 10
@@ -2580,7 +2705,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("SpaceToDepth", ["x"], ["z"], blocksize=b)],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, None, 10, 10))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, None, 10, 10))]
+        )
 
     def test_depth_to_space(self) -> None:
         b = 10
@@ -2806,12 +2933,14 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_value_info("y", TensorProto.FLOAT, (None, None, None, None)),  # type: ignore
+                make_tensor_value_info(
+                    "y", TensorProto.FLOAT, (None, None, None, None)
+                ),
                 make_tensor_value_info(
                     "z", TensorProto.INT64, (None, None, None, None)
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_gemm(self) -> None:
         graph = self._make_graph(
@@ -3177,10 +3306,10 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_value_info("y", TensorProto.FLOAT, (2, None, "b")),  # type: ignore
+                make_tensor_value_info("y", TensorProto.FLOAT, (2, None, "b")),
                 make_tensor_value_info("z", TensorProto.FLOAT, (2, None, "b")),
             ],
-        )  # type: ignore
+        )
 
     def test_split_from_GLU(self) -> None:
         graph = self._make_graph(
@@ -4409,7 +4538,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [],
         )
 
-        self._assert_inferred(graph, [make_tensor_value_info("if_output", TensorProto.FLOAT, (None,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("if_output", TensorProto.FLOAT, (None,))]
+        )
 
     def test_if_no_shape_in_then_branch(self) -> None:
         then_graph = parse_graph(
@@ -4435,7 +4566,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("if_output", TensorProto.FLOAT, None)])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("if_output", TensorProto.FLOAT, None)]
+        )
 
     def test_if_no_shape_in_else_branch(self) -> None:
         then_graph = parse_graph(
@@ -4461,7 +4594,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("if_output", TensorProto.FLOAT, None)])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("if_output", TensorProto.FLOAT, None)]
+        )
 
     def test_if_with_different_optional_shapes_in_then_else_branches(self) -> None:
         # Create a simple If node where the 'then' subgraph adds to the current value, and the 'else' subgraph
@@ -4529,7 +4664,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         output_optional_vi = helper.make_value_info(
             "if_output", output_optional_type_proto
         )
-        self._assert_inferred(graph, [output_optional_vi])  # type: ignore
+        self._assert_inferred(graph, [output_optional_vi])
 
     def test_maxunpool_shape_without_output_shape(self) -> None:
         graph = self._make_graph(
@@ -4580,7 +4715,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("OneHot", ["indices", "depth", "values"], "Y")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (2, 2, None))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (2, 2, None))]
+        )
 
     def test_onehot_with_axis(self) -> None:
         graph = self._make_graph(
@@ -4592,7 +4729,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("OneHot", ["indices", "depth", "values"], "Y", axis=1)],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (2, None, 3, 5))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (2, None, 3, 5))]
+        )
 
     def test_onehot_without_axis_2(self) -> None:
         graph = self._make_graph(
@@ -4605,7 +4744,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [],
             initializer=[make_tensor("depth", TensorProto.INT64, (), (256,))],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (2, 2, 256))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (2, 2, 256))]
+        )
 
     def test_onehot_with_axis_2(self) -> None:
         graph = self._make_graph(
@@ -4618,7 +4759,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [],
             initializer=[make_tensor("depth", TensorProto.INT64, (1,), (256,))],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (2, 256, 3, 5))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (2, 256, 3, 5))]
+        )
 
     def test_loop(self) -> None:
         # can't use self._make_graph for the subgraph as it add more inputs for the Reshape operations it inserts.
@@ -4672,7 +4815,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 ),  # shape may change between iterations
                 make_tensor_value_info("loop_output", TensorProto.FLOAT, (None, 3)),
             ],
-        )  # type: ignore
+        )
 
     def test_loop_no_state(self) -> None:
         input_value_infos = [
@@ -4713,7 +4856,7 @@ class TestShapeInference(TestShapeInferenceHelper):
 
         self._assert_inferred(
             graph, [make_tensor_value_info("loop_output", TensorProto.FLOAT, (None, 3))]
-        )  # type: ignore
+        )
 
     def test_constantofshape_with_input_shape(self) -> None:
         graph = self._make_graph(
@@ -4740,7 +4883,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("shape", TensorProto.INT64, (3,)),
                 make_tensor_value_info("y", TensorProto.INT32, (3, 4, 5)),
             ],
-        )  # type: ignore
+        )
 
     def test_constantofshape_without_input_shape(self) -> None:
         graph = self._make_graph(
@@ -4757,7 +4900,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("y", TensorProto.UINT8, (None, None, None))]
-        )  # type: ignore
+        )
 
     def test_constantofshape_without_input_shape_scalar(self) -> None:
         graph = self._make_graph(
@@ -4774,7 +4917,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("y", TensorProto.UINT8, ())]
-        )  # type: ignore
+        )
 
     def test_constantofshape_with_shape_zero(self) -> None:
         graph = self._make_graph(
@@ -4801,7 +4944,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("shape", TensorProto.INT64, (1,)),
                 make_tensor_value_info("y", TensorProto.INT32, (0,)),
             ],
-        )  # type: ignore
+        )
 
     def test_convinteger(self) -> None:
         graph = self._make_graph(
@@ -4913,7 +5056,10 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.INT32, (30, 50, None, 6, 6))])  # type: ignore
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info("z", TensorProto.INT32, (30, 50, None, 6, 6))],
+        )
 
     def test_convineteger_partial_missing_weight_shape(self) -> None:
         graph = self._make_graph(
@@ -5136,7 +5282,10 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.UINT8, (30, 50, None, 6, 6))])  # type: ignore
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info("y", TensorProto.UINT8, (30, 50, None, 6, 6))],
+        )
 
     def test_qlinearconv_partial_missing_weight_shape(self) -> None:
         graph = self._make_graph(
@@ -5484,12 +5633,12 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_value_info("Y", TensorProto.FLOAT, (None,)),  # type: ignore
-                make_tensor_value_info("indices", TensorProto.INT64, (None,)),  # type: ignore
-                make_tensor_value_info("inverse_indices", TensorProto.INT64, (None,)),  # type: ignore
+                make_tensor_value_info("Y", TensorProto.FLOAT, (None,)),
+                make_tensor_value_info("indices", TensorProto.INT64, (None,)),
+                make_tensor_value_info("inverse_indices", TensorProto.INT64, (None,)),
                 make_tensor_value_info("counts", TensorProto.INT64, (None,)),
             ],
-        )  # type: ignore
+        )
 
     def test_unique_with_axis(self) -> None:
         graph = self._make_graph(
@@ -5507,12 +5656,12 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_value_info("Y", TensorProto.FLOAT, (2, None, 2)),  # type: ignore
-                make_tensor_value_info("indices", TensorProto.INT64, (None,)),  # type: ignore
-                make_tensor_value_info("inverse_indices", TensorProto.INT64, (None,)),  # type: ignore
+                make_tensor_value_info("Y", TensorProto.FLOAT, (2, None, 2)),
+                make_tensor_value_info("indices", TensorProto.INT64, (None,)),
+                make_tensor_value_info("inverse_indices", TensorProto.INT64, (None,)),
                 make_tensor_value_info("counts", TensorProto.INT64, (None,)),
             ],
-        )  # type: ignore
+        )
 
     def test_det(self) -> None:
         graph = self._make_graph(
@@ -5567,7 +5716,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Tile", ["x", "repeats"], ["y"])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (None, None, None))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (None, None, None))]
+        )
 
     @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
     def test_linearclassifier_1D_input(self) -> None:
@@ -5645,7 +5796,10 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, ("num_rois", "C", 10, 5))])  # type: ignore
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info("y", TensorProto.FLOAT, ("num_rois", "C", 10, 5))],
+        )
 
     def test_roialign_symbolic_defaults(self) -> None:
         graph = self._make_graph(
@@ -5657,7 +5811,10 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("RoiAlign", ["x", "rois", "batch_indices"], ["y"])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, ("num_rois", "C", 1, 1))])  # type: ignore
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info("y", TensorProto.FLOAT, ("num_rois", "C", 1, 1))],
+        )
 
     def test_roialign_num_rois(self) -> None:
         graph = self._make_graph(
@@ -5669,7 +5826,62 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("RoiAlign", ["x", "rois", "batch_indices"], ["y"])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (15, "C", 1, 1))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (15, "C", 1, 1))]
+        )
+
+    def test_rotaryembedding_4d(self) -> None:
+        graph = self._make_graph(
+            [
+                ("X", TensorProto.FLOAT, ("B", "num_heads", "seq_len", "head_size")),
+                ("cos_cache", TensorProto.FLOAT, ("max_seq_len", "head_size_div_2")),
+                ("sin_cache", TensorProto.FLOAT, ("max_seq_len", "head_size_div_2")),
+                ("position_ids", TensorProto.INT64, ("B", "seq_len")),
+            ],
+            [
+                make_node(
+                    "RotaryEmbedding",
+                    ["X", "cos_cache", "sin_cache", "position_ids"],
+                    ["Y"],
+                )
+            ],
+            [],
+        )
+        self._assert_inferred(
+            graph,
+            [
+                make_tensor_value_info(
+                    "Y", TensorProto.FLOAT, ("B", "num_heads", "seq_len", "head_size")
+                )
+            ],
+        )
+
+    def test_rotaryembedding_3d(self) -> None:
+        graph = self._make_graph(
+            [
+                ("X", TensorProto.FLOAT, ("B", "seq_len", "hidden_size")),
+                ("cos_cache", TensorProto.FLOAT, ("max_seq_len", "head_size_div_2")),
+                ("sin_cache", TensorProto.FLOAT, ("max_seq_len", "head_size_div_2")),
+                ("position_ids", TensorProto.INT64, ("B", "seq_len")),
+            ],
+            [
+                make_node(
+                    "RotaryEmbedding",
+                    ["X", "cos_cache", "sin_cache", "position_ids"],
+                    ["Y"],
+                    num_heads=4,
+                )
+            ],
+            [],
+        )
+        self._assert_inferred(
+            graph,
+            [
+                make_tensor_value_info(
+                    "Y", TensorProto.FLOAT, ("B", "seq_len", "hidden_size")
+                )
+            ],
+        )
 
     @parameterized.expand(
         all_versions_for("LabelEncoder") if ONNX_ML else [], skip_on_empty=True
@@ -5889,7 +6101,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     keys_tensor=key_tensor,
                     values_tensor=values_tensor,
                     default_tensor=make_tensor(
-                        "default_tensor", TensorProto.STRING, [1], [0]
+                        "default_tensor", TensorProto.INT32, [1], [0]
                     ),
                 )
             ],
@@ -5915,7 +6127,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     keys_tensor=key_tensor,
                     values_strings=["a", "b", "cc", "ddd"],
                     default_tensor=make_tensor(
-                        "default_tensor", TensorProto.STRING, [1, 2], [0, 0]
+                        "default_tensor", TensorProto.STRING, [1, 2], ["a", "b"]
                     ),
                 )
             ],
@@ -5953,7 +6165,9 @@ class TestShapeInference(TestShapeInferenceHelper):
         graph = self._make_graph(
             [], [make_node("Constant", [], ["y"], sparse_value=y_value)], []
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.INT64, y_shape)])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.INT64, y_shape)]
+        )
 
     def test_constant_value_int(self) -> None:
         graph = self._make_graph(
@@ -6040,7 +6254,9 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor("limit", TensorProto.INT32, (), (5,)),
             ],
         )  # Missing 'delta' initializer
-        self._assert_inferred(graph, [make_tensor_value_info("output", TensorProto.INT32, (None,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("output", TensorProto.INT32, (None,))]
+        )
 
     def test_gathernd(self) -> None:
         graph = self._make_graph(
@@ -6084,11 +6300,15 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("NonMaxSuppression", ["boxes", "scores"], ["y"])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.INT64, (None, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.INT64, (None, 3))]
+        )
 
     def test_sequence_empty(self) -> None:
         graph = self._make_graph([], [make_node("SequenceEmpty", [], ["output"])], [])
-        self._assert_inferred(graph, [make_tensor_sequence_value_info("output", TensorProto.FLOAT, None)])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_sequence_value_info("output", TensorProto.FLOAT, None)]
+        )
 
     def test_sequence_construct(self) -> None:
         graph = self._make_graph(
@@ -6113,7 +6333,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (2, 3, 4)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_construct_one_input(self) -> None:
         graph = self._make_graph(
@@ -6128,7 +6348,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (2, 3, 4)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_construct_diff_rank(self) -> None:
         graph = self._make_graph(
@@ -6153,7 +6373,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, None
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_construct_diff_dim_size(self) -> None:
         graph = self._make_graph(
@@ -6178,7 +6398,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (2, 3, None)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_insert(self) -> None:
         graph = self._make_graph(
@@ -6208,7 +6428,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (2, 3, 4)
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_insert_diff_rank(self) -> None:
         graph = self._make_graph(
@@ -6238,7 +6458,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, None
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_insert_diff_shape(self) -> None:
         graph = self._make_graph(
@@ -6261,12 +6481,14 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_sequence_value_info("in_sequence", TensorProto.FLOAT, (2, None, 4)),  # type: ignore
+                make_tensor_sequence_value_info(
+                    "in_sequence", TensorProto.FLOAT, (2, None, 4)
+                ),
                 make_tensor_sequence_value_info(
                     "output_sequence", TensorProto.FLOAT, (2, None, None)
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_at(self) -> None:
         graph = self._make_graph(
@@ -6292,7 +6514,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 ),
                 make_tensor_value_info("output", TensorProto.FLOAT, (2, 3, 4)),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_at_unknown_shape(self) -> None:
         graph = self._make_graph(
@@ -6316,7 +6538,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_sequence_value_info("in_sequence", TensorProto.FLOAT, None),
                 make_tensor_value_info("output", TensorProto.FLOAT, None),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_at_unknown_dim_size(self) -> None:
         graph = self._make_graph(
@@ -6337,10 +6559,12 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_sequence_value_info("in_sequence", TensorProto.FLOAT, (2, 3, None)),  # type: ignore
+                make_tensor_sequence_value_info(
+                    "in_sequence", TensorProto.FLOAT, (2, 3, None)
+                ),
                 make_tensor_value_info("output", TensorProto.FLOAT, (2, 3, None)),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_erase(self) -> None:
         graph = self._make_graph(
@@ -6368,7 +6592,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (2, 3, 4)
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_erase_diff_dim_size(self) -> None:
         graph = self._make_graph(
@@ -6389,12 +6613,14 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_sequence_value_info("in_sequence", TensorProto.FLOAT, (2, None, "x")),  # type: ignore
+                make_tensor_sequence_value_info(
+                    "in_sequence", TensorProto.FLOAT, (2, None, "x")
+                ),
                 make_tensor_sequence_value_info(
                     "output_sequence", TensorProto.FLOAT, (2, None, "x")
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_length(self) -> None:
         graph = self._make_graph(
@@ -6419,7 +6645,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 ),
                 make_tensor_value_info("len", TensorProto.INT64, ()),
             ],
-        )  # type: ignore
+        )
 
     def test_split_to_sequence(self) -> None:
         graph = self._make_graph(
@@ -6435,7 +6661,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (3, 4)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_split_to_sequence_scalar(self) -> None:
         graph = self._make_graph(
@@ -6451,7 +6677,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (2, 4)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_split_to_sequence_keepdims(self) -> None:
         graph = self._make_graph(
@@ -6466,7 +6692,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (1, 4)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_split_to_sequence_not_keepdims(self) -> None:
         graph = self._make_graph(
@@ -6481,7 +6707,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (4,)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_split_to_sequence_ignore_keepdims(self) -> None:
         graph = self._make_graph(
@@ -6504,7 +6730,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (3, 4)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_split_to_sequence_axis(self) -> None:
         graph = self._make_graph(
@@ -6519,7 +6745,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (6, 1)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_split_to_sequence_neg_axis(self) -> None:
         graph = self._make_graph(
@@ -6534,7 +6760,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (1, 4)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_split_to_sequence_split_sizes(self) -> None:
         graph = self._make_graph(
@@ -6550,7 +6776,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (None, 4)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_split_to_sequence_non_divisible(self) -> None:
         graph = self._make_graph(
@@ -6566,7 +6792,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "output_sequence", TensorProto.FLOAT, (None, 4)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_concat_from_sequence(self) -> None:
         graph = self._make_graph(
@@ -6591,7 +6817,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 ),
                 make_tensor_value_info("out", TensorProto.FLOAT, (None, 3, "x")),
             ],
-        )  # type: ignore
+        )
 
     def test_concat_from_sequence_unknown_shape(self) -> None:
         graph = self._make_graph(
@@ -6614,7 +6840,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_sequence_value_info("in_sequence", TensorProto.FLOAT, None),
                 make_tensor_value_info("out", TensorProto.FLOAT, None),
             ],
-        )  # type: ignore
+        )
 
     def test_concat_from_sequence_unknown_dim_size(self) -> None:
         graph = self._make_graph(
@@ -6634,10 +6860,12 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_sequence_value_info("in_sequence", TensorProto.FLOAT, (2, None, "x")),  # type: ignore
+                make_tensor_sequence_value_info(
+                    "in_sequence", TensorProto.FLOAT, (2, None, "x")
+                ),
                 make_tensor_value_info("out", TensorProto.FLOAT, (None, None, "x")),
             ],
-        )  # type: ignore
+        )
 
     def test_concat_from_sequence_axis(self) -> None:
         graph = self._make_graph(
@@ -6657,10 +6885,12 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_sequence_value_info("in_sequence", TensorProto.FLOAT, (2, None, "x")),  # type: ignore
+                make_tensor_sequence_value_info(
+                    "in_sequence", TensorProto.FLOAT, (2, None, "x")
+                ),
                 make_tensor_value_info("out", TensorProto.FLOAT, (2, None, None)),
             ],
-        )  # type: ignore
+        )
 
     def test_concat_from_sequence_neg_axis(self) -> None:
         graph = self._make_graph(
@@ -6680,10 +6910,12 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_sequence_value_info("in_sequence", TensorProto.FLOAT, (2, None, "x")),  # type: ignore
+                make_tensor_sequence_value_info(
+                    "in_sequence", TensorProto.FLOAT, (2, None, "x")
+                ),
                 make_tensor_value_info("out", TensorProto.FLOAT, (None, None, "x")),
             ],
-        )  # type: ignore
+        )
 
     def test_concat_from_sequence_new_axis(self) -> None:
         graph = self._make_graph(
@@ -6710,7 +6942,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 ),
                 make_tensor_value_info("out", TensorProto.FLOAT, (2, 3, None, "x")),
             ],
-        )  # type: ignore
+        )
 
     def test_concat_from_sequence_neg_new_axis(self) -> None:
         graph = self._make_graph(
@@ -6737,7 +6969,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 ),
                 make_tensor_value_info("out", TensorProto.FLOAT, (2, 3, "x", None)),
             ],
-        )  # type: ignore
+        )
 
     def test_adagrad(self) -> None:
         graph = self._make_graph(
@@ -6979,7 +7211,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             graph,
             [make_tensor_value_info("y", TensorProto.FLOAT, (3, None, 4))],
             opset_imports=[helper.make_opsetid(ONNX_DOMAIN, 10)],
-        )  # type: ignore
+        )
 
     def test_constant_pad_2d_opset10(self) -> None:
         graph = self._make_graph(
@@ -7023,7 +7255,9 @@ class TestShapeInference(TestShapeInferenceHelper):
                 )
             ],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, None, 4))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, None, 4))]
+        )
 
     def test_gatherelements_basic(self) -> None:
         graph = self._make_graph(
@@ -7040,11 +7274,13 @@ class TestShapeInference(TestShapeInferenceHelper):
             [
                 ("x", TensorProto.FLOAT, (6,)),
                 ("indices", TensorProto.INT64, None),
-            ],  # type: ignore
+            ],
             [make_node("GatherElements", ["x", "indices"], ["y"])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, None)])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, None)]
+        )
 
     def test_einsum_transpose(self) -> None:
         graph = self._make_graph(
@@ -7052,7 +7288,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ij->ji")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (4, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (4, 3))]
+        )
 
     def test_einsum_dot(self) -> None:
         graph = self._make_graph(
@@ -7060,7 +7298,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="i,i->")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())]
+        )
 
     def test_einsum_scalar(self) -> None:
         graph = self._make_graph(
@@ -7068,7 +7308,19 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation=",->")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())]
+        )
+
+    def test_einsum_scalar_invalid_equation(self) -> None:
+        # Test that scalar inputs with incompatible equations fail gracefully
+        # instead of causing segfaults (issue #6981)
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, ())],
+            [make_node("Einsum", ["x"], ["y"], equation="i->i")],
+            [],
+        )
+        self.assertRaises(onnx.shape_inference.InferenceError, self._inferred, graph)
 
     def test_einsum_outer_prod(self) -> None:
         graph = self._make_graph(
@@ -7076,7 +7328,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,ab->ijab")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 5, 7, 9))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 5, 7, 9))]
+        )
 
     def test_einsum_sum_along_dim(self) -> None:
         graph = self._make_graph(
@@ -7084,15 +7338,19 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="i j->i ")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))])  # type: ignore
-
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))]
+        )
+
     def test_einsum_ellipsis(self) -> None:
         graph = self._make_graph(
             [("x", TensorProto.FLOAT, (3, 4, 4))],
             [make_node("Einsum", ["x"], ["y"], equation="... ii ->... i")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 4))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 4))]
+        )
 
     def test_einsum_ellipsis_2(self) -> None:
         graph = self._make_graph(
@@ -7102,7 +7360,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 5))]
-        )  # type: ignore
+        )
 
     def test_einsum_ellipsis_3(self) -> None:
         graph = self._make_graph(
@@ -7112,7 +7370,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 5))]
-        )  # type: ignore
+        )
 
     def test_einsum_ellipsis_broadcast(self) -> None:
         graph = self._make_graph(
@@ -7122,7 +7380,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("z", TensorProto.FLOAT, (32, 3, 5))]
-        )  # type: ignore
+        )
 
     def test_einsum_contraction(self) -> None:
         graph = self._make_graph(
@@ -7136,7 +7394,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [make_tensor_value_info("z", TensorProto.FLOAT, (5, 6, 7, 9, 10))],
-        )  # type: ignore
+        )
 
     def test_einsum_contraction_2(self) -> None:
         graph = self._make_graph(
@@ -7146,7 +7404,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("z", TensorProto.FLOAT, (4, 5))]
-        )  # type: ignore
+        )
 
     def test_einsum_batch_matmul(self) -> None:
         graph = self._make_graph(
@@ -7154,7 +7412,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="bij , b jk-> bik")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (5, 2, 4))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (5, 2, 4))]
+        )
 
     def test_einsum_left_hand_eqn(self) -> None:
         graph = self._make_graph(
@@ -7162,7 +7422,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,kl")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 3, 4))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 3, 4))]
+        )
 
     def test_einsum_incorrect_num_inputs(self) -> None:
         graph = self._make_graph(
@@ -7182,7 +7444,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="i")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))]
+        )
 
     def test_einsum_sum_A1(self) -> None:  # sums the values of A1
         graph = self._make_graph(
@@ -7190,7 +7454,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="i->")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, ())]
+        )
 
     def test_einsum_element_wise_multiplication_A1_B1(
         self,
@@ -7200,7 +7466,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="i,i->i")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3,))]
+        )
 
     def test_einsum_inner_product_A1_B1(self) -> None:  # inner product of A1 and B1
         graph = self._make_graph(
@@ -7208,7 +7476,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="i,i->")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())]
+        )
 
     def test_einsum_outer_product_A1_B1(self) -> None:  # outer product of A1 and B1
         graph = self._make_graph(
@@ -7216,7 +7486,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="i,j->ij")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_view_A2(self) -> None:  # returns a view of A2
         graph = self._make_graph(
@@ -7224,7 +7496,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ij->ij")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_view_A2_2(self) -> None:  # returns a view of A2, another case
         graph = self._make_graph(
@@ -7232,7 +7506,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ij")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_transpose_A2(self) -> None:  # view transpose of A2
         graph = self._make_graph(
@@ -7240,7 +7516,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ji")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_transpose_A2_to_ij(self) -> None:  # view transpose of A2
         graph = self._make_graph(
@@ -7248,7 +7526,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ji->ij")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_diag_A2(self) -> None:  # view main diagonal of A2
         graph = self._make_graph(
@@ -7256,7 +7536,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ii->i")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))]
+        )
 
     def test_einsum_trace_A2(self) -> None:  # sums main diagonal of A2
         graph = self._make_graph(
@@ -7264,7 +7546,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ii->")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, ())]
+        )
 
     def test_einsum_sum_A2(self) -> None:  # sums the values of A2
         graph = self._make_graph(
@@ -7272,7 +7556,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ij->")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, ())]
+        )
 
     def test_einsum_sum_columns_A2(
         self,
@@ -7282,7 +7568,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ij->j")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))]
+        )
 
     def test_einsum_sum_rows_A2(self) -> None:  # sum horizontally along the rows of A2
         graph = self._make_graph(
@@ -7290,7 +7578,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ij->i")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))]
+        )
 
     def test_einsum_element_wise_multiplication_A2_B2(
         self,
@@ -7300,7 +7590,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,ij->ij")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_element_wise_multiplication_A2_B2_transpose(
         self,
@@ -7310,7 +7602,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,ji->ij")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_matrix_multiplication_A2_B2(
         self,
@@ -7320,7 +7614,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,jk")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_matrix_multiplication_A2_B2_to_ik(
         self,
@@ -7330,7 +7626,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,jk->ik")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_matrix_multiplication_A3_B3(
         self,
@@ -7340,7 +7638,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="bij,bjk->bik")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 3))]
+        )
 
     def test_einsum_matrix_multiplication_A3_B3_transpose(
         self,
@@ -7350,7 +7650,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="bij,bkj->bik")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 3))]
+        )
 
     def test_einsum_inner_product_A2_B2(self) -> None:  # inner product of A2 and B2
         graph = self._make_graph(
@@ -7358,7 +7660,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,kj->ik")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_row_multiplication_A2_B2(
         self,
@@ -7368,7 +7672,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,kj->ikj")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3, 3))]
+        )
 
     def test_einsum_value_multiplication_A2_B2(
         self,
@@ -7378,7 +7684,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,kl->ijkl")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3, 3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3, 3, 3))]
+        )
 
     def test_einsum_scalar_times_array(self) -> None:  # Scalar times array
         graph = self._make_graph(
@@ -7386,7 +7694,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation=",ij->ij")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3, 3))]
+        )
 
     def test_einsum_matrix_vector_A2_B1(self) -> None:  # Matrix and vector.
         graph = self._make_graph(
@@ -7394,7 +7704,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,j->i")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3,))]
+        )
 
     def test_einsum_diag_multiplication_A2_B2(
         self,
@@ -7404,7 +7716,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ii,ii->i")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3,))]
+        )
 
     def test_einsum_diag_dot_product_A2_B2(self) -> None:  # dot product of diagonals
         graph = self._make_graph(
@@ -7412,7 +7726,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ii,ii->")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())]
+        )
 
     def test_negative_log_likehood_shape_is_NCdd(self) -> None:
         N, C = 3, 4
@@ -7428,7 +7744,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("loss", TensorProto.FLOAT, (N,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("loss", TensorProto.FLOAT, (N,))]
+        )
 
     def test_negative_log_likehood_shape_is_NC_with_weight(self) -> None:
         N, C = 3, 4
@@ -7448,7 +7766,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("loss", TensorProto.FLOAT, (N,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("loss", TensorProto.FLOAT, (N,))]
+        )
 
     def test_negative_log_likehood_shape_is_NC_reduction_mean(self) -> None:
         N, C = 3, 4
@@ -7464,7 +7784,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("loss", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("loss", TensorProto.FLOAT, ())]
+        )
 
     def test_negative_log_likehood_shape_is_NC_with_weight_reduction_mean(self) -> None:
         N, C = 3, 4
@@ -7484,7 +7806,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("loss", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("loss", TensorProto.FLOAT, ())]
+        )
 
     def test_negative_log_likehood_shape_is_NCd1d2(self) -> None:
         N, C, d1, d2 = 3, 4, 5, 6
@@ -7503,7 +7827,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("loss", TensorProto.FLOAT, (N, d1, d2))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("loss", TensorProto.FLOAT, (N, d1, d2))]
+        )
 
     def test_negative_log_likehood_shape_is_NCd1d2_with_weight(self) -> None:
         N, C, d1, d2 = 3, 4, 5, 6
@@ -7523,7 +7849,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("loss", TensorProto.FLOAT, (N, d1, d2))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("loss", TensorProto.FLOAT, (N, d1, d2))]
+        )
 
     def test_negative_log_likehood_shape_is_NCd1d2_reduction_sum(self) -> None:
         N, C, d1, d2 = 3, 4, 5, 6
@@ -7542,7 +7870,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("loss", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("loss", TensorProto.FLOAT, ())]
+        )
 
     def test_negative_log_likehood_shape_is_NCd1d2_with_weight_reduction_mean(
         self,
@@ -7564,7 +7894,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("loss", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("loss", TensorProto.FLOAT, ())]
+        )
 
     def test_negative_log_likehood_input_target_shape_mismatch(self) -> None:
         N, C, d1, d2 = 3, 4, 5, 6
@@ -7614,7 +7946,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("SoftmaxCrossEntropyLoss", ["x", "y"], ["z"], reduction="none")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2,))]
+        )
 
     def test_softmax_cross_entropy_mean(self) -> None:
         graph = self._make_graph(
@@ -7622,7 +7956,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("SoftmaxCrossEntropyLoss", ["x", "y"], ["z"], reduction="mean")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())]
+        )
 
     def test_softmax_cross_entropy_none_NCD1D2(self) -> None:
         graph = self._make_graph(
@@ -7633,7 +7969,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("SoftmaxCrossEntropyLoss", ["x", "y"], ["z"], reduction="none")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 5, 8))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 5, 8))]
+        )
 
     def test_softmax_cross_entropy_mean_NCD1D2(self) -> None:
         graph = self._make_graph(
@@ -7644,7 +7982,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("SoftmaxCrossEntropyLoss", ["x", "y"], ["z"], reduction="mean")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, ())]
+        )
 
     def test_celu_function_output_shape(self) -> None:
         graph = self._make_graph(
@@ -7656,24 +7996,24 @@ class TestShapeInference(TestShapeInferenceHelper):
             graph, [make_tensor_value_info("Y", TensorProto.FLOAT, (25, 48, 16, 16))]
         )
 
-    def prepare_input_initializer_tensors(self, initializer_shape, input_shape):  # type: ignore
+    def prepare_input_initializer_tensors(self, initializer_shape, input_shape):
         nodes = [make_node("Add", ["x", "y"], "z")]
         if initializer_shape is None:
-            initializer = []  # type: ignore
+            initializer = []
         else:
             size = 1
             for d in initializer_shape:
                 size = size * d
             vals = [0.0 for i in range(size)]
             initializer = [
-                make_tensor("x", TensorProto.FLOAT, initializer_shape, vals),  # type: ignore
+                make_tensor("x", TensorProto.FLOAT, initializer_shape, vals),
                 make_tensor("y", TensorProto.FLOAT, initializer_shape, vals),
             ]
         if input_shape is None:
-            inputs = []  # type: ignore
+            inputs = []
         else:
             inputs = [
-                helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape),  # type: ignore
+                helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape),
                 helper.make_tensor_value_info("y", TensorProto.FLOAT, input_shape),
             ]
 
@@ -7743,7 +8083,7 @@ class TestShapeInference(TestShapeInferenceHelper):
 
     def test_infer_initializer_input_consistency_all_none(self) -> None:
         initializer_shape = (8, 7)
-        input_shape = (None, None)  # accepatble
+        input_shape = (None, None)  # acceptable
         original_model = self.prepare_input_initializer_tensors(
             initializer_shape, input_shape
         )
@@ -7752,7 +8092,7 @@ class TestShapeInference(TestShapeInferenceHelper):
 
     def test_infer_initializer_input_consistency_single_none(self) -> None:
         initializer_shape = (8, 7)
-        input_shape = (None, 7)  # accepatble
+        input_shape = (None, 7)  # acceptable
         original_model = self.prepare_input_initializer_tensors(
             initializer_shape, input_shape
         )
@@ -7761,7 +8101,7 @@ class TestShapeInference(TestShapeInferenceHelper):
 
     def test_infer_initializer_input_consistency_differnt_rank(self) -> None:
         initializer_shape = (8, 7, 9)
-        input_shape = (None, 7)  # accepatble
+        input_shape = (None, 7)  # acceptable
         original_model = self.prepare_input_initializer_tensors(
             initializer_shape, input_shape
         )
@@ -7777,7 +8117,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         # Reuse test_infer_initializer_input_consistency_all_none test case and check with
         # Serialized model
         initializer_shape = (8, 7)
-        input_shape = (None, None)  # accepatble
+        input_shape = (None, None)  # acceptable
         original_model = self.prepare_input_initializer_tensors(
             initializer_shape, input_shape
         )
@@ -7793,7 +8133,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [],
             initializer=[make_tensor("k", TensorProto.INT64, (), (2,))],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 4, 5))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 4, 5))]
+        )
 
     def test_trilu_lower(self) -> None:
         graph = self._make_graph(
@@ -7802,7 +8144,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [],
             initializer=[make_tensor("k", TensorProto.INT64, (), (10,))],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 4, 5))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3, 4, 5))]
+        )
 
     def test_trilu_upper_zero(self) -> None:
         graph = self._make_graph(
@@ -7811,7 +8155,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [],
             initializer=[make_tensor("k", TensorProto.INT64, (), (5,))],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.INT64, (0, 5))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.INT64, (0, 5))]
+        )
 
     def test_trilu_lower_one(self) -> None:
         graph = self._make_graph(
@@ -7819,7 +8165,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Trilu", ["x"], ["y"], upper=0)],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.INT32, (3, 1, 5))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.INT32, (3, 1, 5))]
+        )
 
     def test_batch_norm_train(self) -> None:
         graph = self._make_graph(
@@ -7843,9 +8191,9 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_value_info("out", TensorProto.FLOAT, (3, 4, 5, 6, 7)),  # type: ignore
-                make_tensor_value_info("output_mean", TensorProto.FLOAT, (4,)),  # type: ignore
-                make_tensor_value_info("output_var", TensorProto.FLOAT, (4,)),  # type: ignore
+                make_tensor_value_info("out", TensorProto.FLOAT, (3, 4, 5, 6, 7)),
+                make_tensor_value_info("output_mean", TensorProto.FLOAT, (4,)),
+                make_tensor_value_info("output_var", TensorProto.FLOAT, (4,)),
             ],
         )
 
@@ -7871,9 +8219,9 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_value_info("out", TensorProto.FLOAT, (3, "C", 5, 6, 7)),  # type: ignore
-                make_tensor_value_info("output_mean", TensorProto.FLOAT, ("C",)),  # type: ignore
-                make_tensor_value_info("output_var", TensorProto.FLOAT, ("C",)),  # type: ignore
+                make_tensor_value_info("out", TensorProto.FLOAT, (3, "C", 5, 6, 7)),
+                make_tensor_value_info("output_mean", TensorProto.FLOAT, ("C",)),
+                make_tensor_value_info("output_var", TensorProto.FLOAT, ("C",)),
             ],
         )
 
@@ -7899,9 +8247,9 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_value_info("out", TensorProto.FLOAT16, (3, 4, 5, 6, 7)),  # type: ignore
-                make_tensor_value_info("output_mean", TensorProto.FLOAT, (4,)),  # type: ignore
-                make_tensor_value_info("output_var", TensorProto.FLOAT, (4,)),  # type: ignore
+                make_tensor_value_info("out", TensorProto.FLOAT16, (3, 4, 5, 6, 7)),
+                make_tensor_value_info("output_mean", TensorProto.FLOAT, (4,)),
+                make_tensor_value_info("output_var", TensorProto.FLOAT, (4,)),
             ],
         )
 
@@ -7924,7 +8272,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("out", TensorProto.FLOAT, (3, 4, 5, 6, 7))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("out", TensorProto.FLOAT, (3, 4, 5, 6, 7))]
+        )
 
     def test_batch_norm_test_no_dim(self) -> None:
         graph = self._make_graph(
@@ -7945,7 +8295,14 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("out", TensorProto.FLOAT, (3, 4, None, None, None))])  # type: ignore
+        self._assert_inferred(
+            graph,
+            [
+                make_tensor_value_info(
+                    "out", TensorProto.FLOAT, (3, 4, None, None, None)
+                )
+            ],
+        )
 
     def test_batch_norm_train_no_shape(self) -> None:
         graph = self._make_graph(
@@ -7969,9 +8326,9 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [
-                make_tensor_value_info("out", TensorProto.FLOAT, None),  # type: ignore
-                make_tensor_value_info("running_mean", TensorProto.FLOAT, ("C",)),  # type: ignore
-                make_tensor_value_info("running_var", TensorProto.FLOAT, ("C",)),  # type: ignore
+                make_tensor_value_info("out", TensorProto.FLOAT, None),
+                make_tensor_value_info("running_mean", TensorProto.FLOAT, ("C",)),
+                make_tensor_value_info("running_var", TensorProto.FLOAT, ("C",)),
             ],
         )
 
@@ -7981,13 +8338,17 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("NonZero", ["x"], ["out"])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("out", TensorProto.INT64, (1, None))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("out", TensorProto.INT64, (1, None))]
+        )
 
     def test_nonzero_no_shape(self) -> None:
         graph = self._make_graph(
             [("x", TensorProto.FLOAT, None)], [make_node("NonZero", ["x"], ["out"])], []
         )
-        self._assert_inferred(graph, [make_tensor_value_info("out", TensorProto.INT64, (None, None))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("out", TensorProto.INT64, (None, None))]
+        )
 
     def test_nonzero_existing_dim_param(self) -> None:
         graph = self._make_graph(
@@ -7995,13 +8356,17 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("NonZero", ["x"], ["y"])],
             [make_tensor_value_info("y", TensorProto.INT64, (None, "NZ"))],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.INT64, (1, "NZ"))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("y", TensorProto.INT64, (1, "NZ"))]
+        )
 
     def test_nonzero_scalar(self) -> None:
         graph = self._make_graph(
             [("x", TensorProto.FLOAT, ())], [make_node("NonZero", ["x"], ["out"])], []
         )
-        self._assert_inferred(graph, [make_tensor_value_info("out", TensorProto.INT64, (0, None))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("out", TensorProto.INT64, (0, None))]
+        )
 
     def test_optional_construct_empty_tensor(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(
@@ -8014,7 +8379,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         graph = self._make_graph(
             [], [make_node("Optional", [], ["output"], type=tensor_type_proto)], []
         )
-        self._assert_inferred(graph, [optional_val_info])  # type: ignore
+        self._assert_inferred(graph, [optional_val_info])
 
     def test_optional_construct_empty_sequence(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(
@@ -8030,7 +8395,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Optional", [], ["output_sequence"], type=sequence_type_proto)],
             [],
         )
-        self._assert_inferred(graph, [optional_val_info])  # type: ignore
+        self._assert_inferred(graph, [optional_val_info])
 
     def test_optional_construct_tensor(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(
@@ -8045,7 +8410,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Optional", ["input1"], ["output"])],
             [],
         )
-        self._assert_inferred(graph, [optional_val_info])  # type: ignore
+        self._assert_inferred(graph, [optional_val_info])
 
     def test_optional_construct_sequence(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(
@@ -8067,7 +8432,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [sequence_val_info, optional_val_info])  # type: ignore
+        self._assert_inferred(graph, [sequence_val_info, optional_val_info])
 
     def test_optional_tensor_has_element(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(
@@ -8088,7 +8453,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [optional_val_info, make_tensor_value_info("output", TensorProto.BOOL, ())],
-        )  # type: ignore
+        )
 
     def test_optional_sequence_has_element(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(
@@ -8118,7 +8483,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 optional_val_info,
                 make_tensor_value_info("output", TensorProto.BOOL, ()),
             ],
-        )  # type: ignore
+        )
 
     def test_tensor_get_element(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(
@@ -8134,7 +8499,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [output_tensor_val_info])  # type: ignore
+        self._assert_inferred(graph, [output_tensor_val_info])
 
     @parameterized.expand(all_versions_for("StringSplit"))
     def test_string_split_basic(self, _, version) -> None:
@@ -8254,7 +8619,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [optional_val_info, tensor_val_into])  # type: ignore
+        self._assert_inferred(graph, [optional_val_info, tensor_val_into])
 
     def test_optional_sequence_get_element(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(
@@ -8280,7 +8645,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
             [],
         )
-        self._assert_inferred(graph, [optional_val_info, sequence_val_into, output_val_into])  # type: ignore
+        self._assert_inferred(
+            graph, [optional_val_info, sequence_val_into, output_val_into]
+        )
 
     def test_where_bfloat(self) -> None:
         graph = self._make_graph(
@@ -8292,7 +8659,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Where", ["cond", "x", "y"], ["out"])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("out", TensorProto.BFLOAT16, (10,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("out", TensorProto.BFLOAT16, (10,))]
+        )
 
     def test_parse_data_with_unsupported_tensor_type(self) -> None:
         model = helper.make_model(
@@ -8381,7 +8750,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("grid", TensorProto.FLOAT, (N, H, W, 2))]
-        )  # type: ignore
+        )
 
     def test_affinegrid_3d(self) -> None:
         N, C, D, H, W = 2, 3, 4, 5, 6
@@ -8402,7 +8771,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("grid", TensorProto.FLOAT, (N, D, H, W, 3))]
-        )  # type: ignore
+        )
 
     def test_gridsample_2d(self) -> None:
         graph = self._make_graph(
@@ -8424,7 +8793,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("y", TensorProto.FLOAT, (1, 1, 3, 3))]
-        )  # type: ignore
+        )
 
     def test_gridsample_3d(self) -> None:
         graph = self._make_graph(
@@ -8446,7 +8815,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph, [make_tensor_value_info("y", TensorProto.FLOAT, (1, 1, 3, 2, 3))]
-        )  # type: ignore
+        )
 
     def test_gridsample_2d_defaults(self) -> None:
         graph = self._make_graph(
@@ -8464,7 +8833,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "y", TensorProto.FLOAT, ("N", "C", "H_out", "W_out")
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_gridsample_3d_defaults(self) -> None:
         graph = self._make_graph(
@@ -8482,7 +8851,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "y", TensorProto.FLOAT, ("N", "C", "D_out", "H_out", "W_out")
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_gridsample_2d_no_dim(self) -> None:
         graph = self._make_graph(
@@ -8504,7 +8873,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         self._assert_inferred(
             graph,
             [make_tensor_value_info("y", TensorProto.FLOAT, ("N", "C", None, None))],
-        )  # type: ignore
+        )
 
     def test_gridsample_3d_no_dim(self) -> None:
         graph = self._make_graph(
@@ -8530,7 +8899,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "y", TensorProto.FLOAT, ("N", "C", None, None, None)
                 )
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_map_identity_known_dims(self):
         input_value_infos = [
@@ -8571,7 +8940,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "out_sequence", TensorProto.FLOAT, (220, 220, 3)
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_map_identity_unknown_dims(self):
         input_value_infos = [
@@ -8612,7 +8981,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "out_sequence", TensorProto.FLOAT, (None, None, 3)
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_map_slice_outs_known_dims(self):
         body_graph = helper.make_graph(
@@ -8641,7 +9010,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor("starts2", TensorProto.INT64, (2,), (0, 0)),
                 make_tensor("ends2", TensorProto.INT64, (2,), (30, 40)),
             ],
-        )  # type: ignore
+        )
 
         graph = self._make_graph(
             [
@@ -8675,7 +9044,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "out_sequence2", TensorProto.FLOAT, (30, 40, 3)
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_map_slice_outs_unknown_dims(self):
         body_graph = helper.make_graph(
@@ -8704,7 +9073,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor("starts2", TensorProto.INT64, (2,), (0, 0)),
                 make_tensor("ends2", TensorProto.INT64, (2,), (30, 40)),
             ],
-        )  # type: ignore
+        )
 
         graph = self._make_graph(
             [
@@ -8738,7 +9107,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "out_sequence2", TensorProto.FLOAT, (None, None, 3)
                 ),
             ],
-        )  # type: ignore
+        )
 
     def test_sequence_map_different_tensor_type(self):
         body_graph = helper.make_graph(
@@ -8754,7 +9123,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                     "shape", onnx.TensorProto.INT64, (3,)
                 )
             ],
-        )  # type: ignore
+        )
 
         graph = self._make_graph(
             [
@@ -8778,7 +9147,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 ),
                 make_tensor_sequence_value_info("shapes", TensorProto.INT64, (3,)),
             ],
-        )  # type: ignore
+        )
 
     def test_hammingwindow(self):
         graph = self._make_graph(
@@ -8800,7 +9169,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("shape", TensorProto.INT64, ()),
                 make_tensor_value_info("y", TensorProto.FLOAT, (10,)),
             ],
-        )  # type: ignore
+        )
 
         graph = self._make_graph(
             [],
@@ -8821,7 +9190,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("shape", TensorProto.INT64, ()),
                 make_tensor_value_info("y", TensorProto.FLOAT, (10,)),
             ],
-        )  # type: ignore
+        )
 
     def test_hannwindow(self):
         graph = self._make_graph(
@@ -8843,7 +9212,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("shape", TensorProto.INT64, ()),
                 make_tensor_value_info("y", TensorProto.FLOAT, (10,)),
             ],
-        )  # type: ignore
+        )
 
         graph = self._make_graph(
             [],
@@ -8864,7 +9233,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("shape", TensorProto.INT64, ()),
                 make_tensor_value_info("y", TensorProto.FLOAT, (10,)),
             ],
-        )  # type: ignore
+        )
 
     def test_blackmanwindow(self):
         graph = self._make_graph(
@@ -8886,7 +9255,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("shape", TensorProto.INT64, ()),
                 make_tensor_value_info("y", TensorProto.FLOAT, (10,)),
             ],
-        )  # type: ignore
+        )
 
         graph = self._make_graph(
             [],
@@ -8907,7 +9276,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("shape", TensorProto.INT64, ()),
                 make_tensor_value_info("y", TensorProto.FLOAT, (10,)),
             ],
-        )  # type: ignore
+        )
 
     @parameterized.expand(
         [
@@ -8969,7 +9338,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         if version < 20:
             if axis is not None:
                 attributes["axis"] = axis
-            nodes = [make_node("DFT", ["input", ""], ["output"], **attributes)]  # type: ignore[arg-type]
+            nodes = [make_node("DFT", ["input", ""], ["output"], **attributes)]
             value_infos = []
         else:
             assert version >= 20
@@ -8981,12 +9350,12 @@ class TestShapeInference(TestShapeInferenceHelper):
                         ["axis"],
                         value=make_tensor("axis", TensorProto.INT64, (), (axis,)),
                     ),
-                    make_node("DFT", ["input", "", "axis"], ["output"], **attributes),  # type: ignore[arg-type]
+                    make_node("DFT", ["input", "", "axis"], ["output"], **attributes),
                 ]
                 value_infos = [make_tensor_value_info("axis", TensorProto.INT64, ())]
             else:
                 nodes = [
-                    make_node("DFT", ["input", "", ""], ["output"], **attributes),  # type: ignore[arg-type]
+                    make_node("DFT", ["input", "", ""], ["output"], **attributes),
                 ]
                 value_infos = []
 
@@ -9090,7 +9459,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                         "dft_length", TensorProto.INT64, (), (dft_length,)
                     ),
                 ),
-                make_node("DFT", ["input", "dft_length"], ["output"], **attributes),  # type: ignore[arg-type]
+                make_node("DFT", ["input", "dft_length"], ["output"], **attributes),
             ]
             value_infos = [make_tensor_value_info("dft_length", TensorProto.INT64, ())]
         else:
@@ -9111,7 +9480,12 @@ class TestShapeInference(TestShapeInferenceHelper):
                             "dft_length", TensorProto.INT64, (), (dft_length,)
                         ),
                     ),
-                    make_node("DFT", ["input", "dft_length", "axis"], ["output"], **attributes),  # type: ignore[arg-type]
+                    make_node(
+                        "DFT",
+                        ["input", "dft_length", "axis"],
+                        ["output"],
+                        **attributes,
+                    ),
                 ]
                 value_infos = [
                     make_tensor_value_info("dft_length", TensorProto.INT64, ()),
@@ -9127,7 +9501,12 @@ class TestShapeInference(TestShapeInferenceHelper):
                             "dft_length", TensorProto.INT64, (), (dft_length,)
                         ),
                     ),
-                    make_node("DFT", ["input", "dft_length", ""], ["output"], **attributes),  # type: ignore[arg-type]
+                    make_node(
+                        "DFT",
+                        ["input", "dft_length", ""],
+                        ["output"],
+                        **attributes,
+                    ),
                 ]
                 value_infos = [
                     make_tensor_value_info("dft_length", TensorProto.INT64, ())
@@ -9467,7 +9846,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("window", TensorProto.INT64, (5,)),
                 make_tensor_value_info("output", TensorProto.FLOAT, (2, 3, 5, 2)),
             ],
-        )  # type: ignore
+        )
 
         graph = self._make_graph(
             [],
@@ -9517,7 +9896,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("frame_length", TensorProto.INT64, ()),
                 make_tensor_value_info("output", TensorProto.FLOAT, (2, 3, 5, 2)),
             ],
-        )  # type: ignore
+        )
 
         graph = self._make_graph(
             [],
@@ -9560,7 +9939,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("frame_length", TensorProto.INT64, ()),
                 make_tensor_value_info("output", TensorProto.FLOAT, (2, 3, 5, 2)),
             ],
-        )  # type: ignore
+        )
 
     def test_melweightmatrix(self):
         graph = self._make_graph(
@@ -9624,7 +10003,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("upper_edge_hertz", TensorProto.FLOAT, ()),
                 make_tensor_value_info("output", TensorProto.FLOAT, (65, 10)),
             ],
-        )  # type: ignore
+        )
 
     def test_melweightmatrix_with_output_datatype(self):
         graph = self._make_graph(
@@ -9689,7 +10068,7 @@ class TestShapeInference(TestShapeInferenceHelper):
                 make_tensor_value_info("upper_edge_hertz", TensorProto.FLOAT, ()),
                 make_tensor_value_info("output", TensorProto.DOUBLE, (65, 10)),
             ],
-        )  # type: ignore
+        )
 
     def test_center_crop_pad_hwc_crop(self):
         graph = self._make_graph(
@@ -9795,6 +10174,8 @@ class TestShapeInference(TestShapeInferenceHelper):
             ["x"],
             ["y"],
             domain=ONNX_ML_DOMAIN,
+            cats_int64s=[1, 2, 3],
+            cats_strings=["1", "2", "3"],
         )
         graph_int = self._make_graph(
             [("x", TensorProto.INT64, (30, 4, 5))],
@@ -9823,6 +10204,41 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
         )
 
+    @parameterized.expand(
+        [
+            ([1, 2, 3], ["1", "2"]),
+            ([1, 2, 3], None),
+            (None, ["1", "2", "3"]),
+            (None, None),
+        ]
+    )
+    @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
+    def test_category_mapper_fails_if_invalid_attributes(
+        self, cats_int64s, cats_strings
+    ) -> None:
+        cat = make_node(
+            "CategoryMapper",
+            ["x"],
+            ["y"],
+            domain=ONNX_ML_DOMAIN,
+            cats_int64s=cats_int64s,
+            cats_strings=cats_strings,
+        )
+        graph = self._make_graph(
+            [("x", TensorProto.INT64, (30, 4, 5))],
+            [cat],
+            [],
+        )
+        self.assertRaises(
+            onnx.shape_inference.InferenceError,
+            self._inferred,
+            graph,
+            opset_imports=[
+                make_opsetid(ONNX_ML_DOMAIN, 1),
+                make_opsetid(ONNX_DOMAIN, 11),
+            ],
+        )
+
     @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
     def test_tree_ensemble_regressor(self) -> None:
         tree = make_node(
@@ -9902,42 +10318,26 @@ class TestShapeInference(TestShapeInferenceHelper):
 
     @parameterized.expand(
         [
-            {
-                "nodes_truenodeids": [0] * 6,
-                "leaf_weights": make_tensor(
-                    "leaf_weights", TensorProto.DOUBLE, (9,), [1] * 9
-                ),
-                "nodes_splits": make_tensor(
-                    "nodes_splits", TensorProto.DOUBLE, (5,), [1] * 5
-                ),
-            },
-            {
-                "nodes_truenodeids": [0] * 5,
-                "leaf_weights": make_tensor(
-                    "leaf_weights", TensorProto.FLOAT, (9,), [1] * 9
-                ),
-                "nodes_splits": make_tensor(
-                    "nodes_splits", TensorProto.DOUBLE, (5,), [1] * 5
-                ),
-            },
-            {
-                "nodes_truenodeids": [0] * 5,
-                "leaf_weights": make_tensor(
-                    "leaf_weights", TensorProto.DOUBLE, (18,), [1] * 18
-                ),
-                "nodes_splits": make_tensor(
-                    "nodes_splits", TensorProto.DOUBLE, (5,), [1] * 5
-                ),
-            },
-            {
-                "nodes_truenodeids": [0] * 5,
-                "leaf_weights": make_tensor(
-                    "leaf_weights", TensorProto.DOUBLE, (9,), [1] * 9
-                ),
-                "nodes_splits": make_tensor(
-                    "nodes_splits", TensorProto.FLOAT, (5,), [1] * 5
-                ),
-            },
+            (
+                [0] * 6,
+                make_tensor("leaf_weights", TensorProto.DOUBLE, (9,), [1] * 9),
+                make_tensor("nodes_splits", TensorProto.DOUBLE, (5,), [1] * 5),
+            ),
+            (
+                [0] * 5,
+                make_tensor("leaf_weights", TensorProto.FLOAT, (9,), [1] * 9),
+                make_tensor("nodes_splits", TensorProto.DOUBLE, (5,), [1] * 5),
+            ),
+            (
+                [0] * 5,
+                make_tensor("leaf_weights", TensorProto.DOUBLE, (18,), [1] * 18),
+                make_tensor("nodes_splits", TensorProto.DOUBLE, (5,), [1] * 5),
+            ),
+            (
+                [0] * 5,
+                make_tensor("leaf_weights", TensorProto.DOUBLE, (9,), [1] * 9),
+                make_tensor("nodes_splits", TensorProto.FLOAT, (5,), [1] * 5),
+            ),
         ]
     )
     @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
@@ -9977,7 +10377,15 @@ class TestShapeInference(TestShapeInferenceHelper):
             [tree],
             [],
         )
-        self.assertRaises(onnx.shape_inference.InferenceError, self._inferred, graph)
+        self.assertRaises(
+            onnx.shape_inference.InferenceError,
+            self._inferred,
+            graph,
+            opset_imports=[
+                make_opsetid(ONNX_ML_DOMAIN, 5),
+                make_opsetid(ONNX_DOMAIN, 11),
+            ],
+        )
 
     @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
     def test_tree_ensemble_classifier(self) -> None:
@@ -10028,7 +10436,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             )
             self._assert_inferred(
                 graph,
-                [make_tensor_value_info("z", TensorProto.INT64, (3, 4, expected))],  # type: ignore
+                [make_tensor_value_info("z", TensorProto.INT64, (3, 4, expected))],
                 opset_imports=[
                     make_opsetid(ONNX_ML_DOMAIN, 3),
                     make_opsetid(ONNX_DOMAIN, 18),
@@ -10052,7 +10460,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph,
-            [make_tensor_value_info("y", TensorProto.INT64, (3, 4, 5))],  # type: ignore
+            [make_tensor_value_info("y", TensorProto.INT64, (3, 4, 5))],
             opset_imports=[
                 make_opsetid(ONNX_ML_DOMAIN, 3),
                 make_opsetid(ONNX_DOMAIN, 18),
@@ -10083,6 +10491,40 @@ class TestShapeInference(TestShapeInferenceHelper):
             ],
         )
 
+    @parameterized.expand(
+        [
+            ([1, 2, 3], ["1", "2", "3"]),
+            (None, None),
+        ]
+    )
+    @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
+    def test_one_hot_encoder_fails_if_invalid_attributes(
+        self, cats_int64s, cats_strings
+    ) -> None:
+        graph = self._make_graph(
+            [("input", TensorProto.INT64, (2, "N", 3))],
+            [
+                make_node(
+                    "OneHotEncoder",
+                    ["input"],
+                    ["output"],
+                    cats_int64s=cats_int64s,
+                    cats_strings=cats_strings,
+                    domain="ai.onnx.ml",
+                )
+            ],
+            [],
+        )
+        self.assertRaises(
+            onnx.shape_inference.InferenceError,
+            self._inferred,
+            graph,
+            opset_imports=[
+                make_opsetid(ONNX_ML_DOMAIN, 1),
+                make_opsetid(ONNX_DOMAIN, 11),
+            ],
+        )
+
     @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
     def test_zip_map(self) -> None:
         params = (
@@ -10132,7 +10574,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Compress", ["input", "condition"], ["output"])],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("output", TensorProto.INT64, (None,))])  # type: ignore
+        self._assert_inferred(
+            graph, [make_tensor_value_info("output", TensorProto.INT64, (None,))]
+        )
 
     def test_compress_with_axis(self) -> None:
         graph = self._make_graph(
@@ -10143,7 +10587,10 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Compress", ["input", "condition"], ["output"], axis=-1)],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("output", TensorProto.INT64, (2, "N", 3, None))])  # type: ignore
+        self._assert_inferred(
+            graph,
+            [make_tensor_value_info("output", TensorProto.INT64, (2, "N", 3, None))],
+        )
 
     def test_check_type_when_schema_has_empty_io(self):
         input = """
@@ -10232,6 +10679,298 @@ class TestShapeInference(TestShapeInferenceHelper):
             onnx.checker.check_model(model, full_check=True)
             onnx.shape_inference.infer_shapes(model)
 
+    @parameterized.expand(all_versions_for("ConstantOfShape"))
+    def test_issue_constantofshape_6135(self, _, version):
+        graph = self._make_graph(
+            [("std.constant", TensorProto.INT64, (1,)), "output"],
+            [
+                make_node(
+                    "ConstantOfShape",
+                    inputs=["std.constant"],
+                    outputs=["output"],
+                    name="invalid_node",
+                )
+            ],
+            [],
+            initializer=[make_tensor("std.constant", TensorProto.FLOAT, (1,), (-10,))],
+        )
+        self.assertRaises(
+            onnx.shape_inference.InferenceError,
+            self._inferred,
+            graph,
+            opset_imports=[helper.make_opsetid(ONNX_DOMAIN, version)],
+        )
+
+
+class TestCustomSchemaShapeInference(TestShapeInferenceHelper):
+    custom_op_type: str = "CustomOp"
+    dummy_graph_op_type: str = "DummyGraph"
+    op_version: int = 1
+    op_domain: str = ""
+
+    def setUp(self) -> None:
+        # Ensure the schema is unregistered
+        self.assertFalse(onnx.defs.has(self.custom_op_type, self.op_domain))
+        self.assertFalse(onnx.defs.has(self.dummy_graph_op_type, self.op_domain))
+
+    def tearDown(self) -> None:
+        # Clean up the registered schema
+        with contextlib.suppress(onnx.defs.SchemaError):
+            onnx.defs.deregister_schema(
+                self.custom_op_type, self.op_version, self.op_domain
+            )
+            onnx.defs.deregister_schema(
+                self.dummy_graph_op_type, self.op_version, self.op_domain
+            )
+
+    def get_custom_op_schema(self):
+        # CustomOp schema:
+        #   attrs:
+        #       out_len: [L0, L1, ...]
+        #   inputs:
+        #       a[N, La]
+        #       b[N, Lb]
+        #   outputs:
+        #       out0[N, La * Lb, L0]
+        #       out1[N, La * Lb, L1]
+        #       ...
+        schema = OpSchema(
+            self.custom_op_type,
+            self.op_domain,
+            self.op_version,
+            inputs=[
+                defs.OpSchema.FormalParameter("a", "float"),
+                defs.OpSchema.FormalParameter("b", "float"),
+            ],
+            outputs=[
+                defs.OpSchema.FormalParameter(
+                    "out", "float", param_option=OpSchema.FormalParameterOption.Variadic
+                ),
+            ],
+            attributes=[
+                defs.OpSchema.Attribute("out_len", defs.OpSchema.AttrType.INTS)
+            ],
+        )
+
+        def schema_shape_infer_func(ctx: onnx.shape_inference.InferenceContext):
+            def parse_tensor_input(t: TypeProto):
+                self.assertTrue(isinstance(t, TypeProto))
+                return (
+                    t.tensor_type.elem_type,
+                    [
+                        d.dim_value if d.HasField("dim_value") else None
+                        for d in t.tensor_type.shape.dim
+                    ],
+                )
+
+            self.assertEqual(ctx.get_num_inputs(), 2)
+            in0 = ctx.get_input_type(0)
+            in1 = ctx.get_input_type(1)
+            in0_type, in0_shape = parse_tensor_input(in0)
+            in1_type, in1_shape = parse_tensor_input(in1)
+            self.assertEqual(in0_type, TensorProto.FLOAT)
+            self.assertEqual(in1_type, TensorProto.FLOAT)
+            self.assertEqual(len(in0_shape), 2)
+            self.assertEqual(len(in1_shape), 2)
+            self.assertEqual(in0_shape[0], in1_shape[0])
+            N, La = in0_shape
+            _, Lb = in1_shape
+            attr = ctx.get_attribute("out_len")
+            out_len = attr.ints
+            self.assertEqual(len(out_len), ctx.get_num_outputs())
+            for i in range(ctx.get_num_outputs()):
+                out = ctx.get_output_type(i)
+                out.tensor_type.elem_type = in0_type
+                out.tensor_type.shape.dim.add().dim_value = N
+                out.tensor_type.shape.dim.add().dim_value = La * Lb
+                out.tensor_type.shape.dim.add().dim_value = out_len[i]
+                ctx.set_output_type(i, out)
+
+        schema.set_type_and_shape_inference_function(schema_shape_infer_func)
+        return schema
+
+    def get_dummy_graph_schema(self):
+        # DummyGraph schema:
+        #   attrs:
+        #       graph: OnnxGraph
+        #   inputs:
+        #       as same as the graph attribute
+        #   outputs:
+        #       as same as the graph attribute
+        schema = OpSchema(
+            self.dummy_graph_op_type,
+            self.op_domain,
+            self.op_version,
+            inputs=[
+                defs.OpSchema.FormalParameter(
+                    "in", "float", param_option=OpSchema.FormalParameterOption.Variadic
+                ),
+            ],
+            outputs=[
+                defs.OpSchema.FormalParameter(
+                    "out", "float", param_option=OpSchema.FormalParameterOption.Variadic
+                ),
+            ],
+            attributes=[defs.OpSchema.Attribute("graph", defs.OpSchema.AttrType.GRAPH)],
+        )
+
+        def schema_shape_infer_func(ctx: onnx.shape_inference.InferenceContext):
+            self.assertEqual(ctx.get_num_inputs(), 2)
+            self.assertIsNotNone(ctx.get_attribute("graph"))
+            gctx = ctx.get_graph_attribute_inferencer("graph")
+            outputs = gctx.do_inferencing(
+                [ctx.get_input_type(i) for i in range(ctx.get_num_inputs())],
+                [ctx.get_input_data(i) for i in range(ctx.get_num_inputs())],
+            )
+            for idx, out in enumerate(outputs):
+                ctx.set_output_type(idx, out)
+
+        schema.set_type_and_shape_inference_function(schema_shape_infer_func)
+        return schema
+
+    def gen_custom_op_graph(self, N, La, Lb, out_len, mark_output=False):
+        a = make_tensor_value_info("a", TensorProto.FLOAT, (N, La))
+        b = make_tensor_value_info("b", TensorProto.FLOAT, (N, Lb))
+        outs = [
+            make_tensor_value_info(f"out{i}", TensorProto.FLOAT, None)
+            for i in range(len(out_len))
+        ]
+        node = make_node(
+            self.custom_op_type, ["a", "b"], [v.name for v in outs], out_len=out_len
+        )
+        graph = make_graph(
+            [node], "test", [a, b], outs if mark_output else [], value_info=outs
+        )
+        return graph
+
+    def gen_dummy_graph_graph(self, N, La, Lb, out_len):
+        subgraph = self.gen_custom_op_graph(N, La, Lb, out_len, True)
+        a = make_tensor_value_info("a", TensorProto.FLOAT, (N, La))
+        b = make_tensor_value_info("b", TensorProto.FLOAT, (N, Lb))
+        outs = [
+            make_tensor_value_info(f"out{i}", TensorProto.FLOAT, None)
+            for i in range(len(out_len))
+        ]
+        node = make_node(
+            self.dummy_graph_op_type, ["a", "b"], [v.name for v in outs], graph=subgraph
+        )
+        graph = make_graph([node], "test", [a, b], [], value_info=outs)
+        return graph
+
+    def shape_infer_once(self, graph, N, La, Lb, out_len):
+        self._assert_inferred(
+            graph,
+            [
+                make_tensor_value_info(f"out{i}", TensorProto.FLOAT, (N, La * Lb, Li))
+                for i, Li in enumerate(out_len)
+            ],
+        )
+
+    def test_custom_schema_shape_inference(self) -> None:
+        # generate graph
+        N = 3
+        La = 32
+        Lb = 64
+        out_len = [1, 2]
+        graph = self.gen_custom_op_graph(N, La, Lb, out_len)
+
+        # shape inference before register
+        with self.assertRaises(onnx.checker.ValidationError):
+            self.shape_infer_once(graph, N, La, Lb, out_len)
+
+        # register schema
+        schema = self.get_custom_op_schema()
+        onnx.defs.register_schema(schema)
+
+        # shape inference with registered schema
+        self.shape_infer_once(graph, N, La, Lb, out_len)
+
+        # clean up
+        onnx.defs.deregister_schema(schema.name, schema.since_version, schema.domain)
+
+    def test_dummy_graph_schema_shape_inference(self) -> None:
+        # generate graph
+        N = 3
+        La = 32
+        Lb = 64
+        out_len = [1, 2]
+        graph = self.gen_dummy_graph_graph(N, La, Lb, out_len)
+
+        # shape inference before register
+        with self.assertRaises(onnx.checker.ValidationError):
+            self.shape_infer_once(graph, N, La, Lb, out_len)
+
+        # register schema
+        custom_op_schema = self.get_custom_op_schema()
+        dummy_graph_schema = self.get_dummy_graph_schema()
+        onnx.defs.register_schema(custom_op_schema)
+        onnx.defs.register_schema(dummy_graph_schema)
+
+        # shape inference with registered schema
+        self.shape_infer_once(graph, N, La, Lb, out_len)
+
+        # clean up
+        onnx.defs.deregister_schema(
+            custom_op_schema.name,
+            custom_op_schema.since_version,
+            custom_op_schema.domain,
+        )
+        onnx.defs.deregister_schema(
+            dummy_graph_schema.name,
+            dummy_graph_schema.since_version,
+            dummy_graph_schema.domain,
+        )
+
+    def test_invalid_field_in_inference_func(self) -> None:
+        N = 3
+        La = 32
+        Lb = 64
+        out_len = [1]
+        graph = self.gen_custom_op_graph(N, La, Lb, out_len)
+
+        schema = self.get_custom_op_schema()
+        raw_func = schema.get_type_and_shape_inference_function()
+
+        def schema_shape_infer_func(ctx: onnx.shape_inference.InferenceContext):
+            raw_func(ctx)
+            self.assertIsNone(ctx.get_attribute("not-exist-attr"))
+            self.assertTrue(ctx.has_input(0))
+            self.assertFalse(ctx.has_input(2))
+            with self.assertRaises(TypeError):
+                self.assertFalse(ctx.has_input(-1))
+            self.assertTrue(ctx.has_output(0))
+            self.assertFalse(ctx.has_output(1))
+            with self.assertRaises(TypeError):
+                self.assertFalse(ctx.has_output(-1))
+            with self.assertRaises(onnx.shape_inference.InferenceError):
+                ctx.get_graph_attribute_inferencer("not-exist-attr")
+            self.assertIsNone(ctx.get_input_data(0))
+            with self.assertRaises(RuntimeError):
+                self.assertIsNone(ctx.get_input_data(10))
+            self.assertIsNone(ctx.get_input_sparse_data(0))
+            with self.assertRaises(RuntimeError):
+                self.assertIsNone(ctx.get_input_sparse_data(10))
+            self.assertIsNotNone(ctx.get_input_type(0))
+            with self.assertRaises(RuntimeError):
+                ctx.get_input_type(10)
+            self.assertIsNone(ctx.get_symbolic_input(0))
+            with self.assertRaises(RuntimeError):
+                ctx.get_symbolic_input(10)
+            self.assertIsNotNone(ctx.get_output_type(0))
+            with self.assertRaises(RuntimeError):
+                ctx.get_output_type(10)
+            self.assertEqual(ctx.get_num_inputs(), 2)
+            self.assertEqual(ctx.get_num_outputs(), 1)
+
+        schema.set_type_and_shape_inference_function(schema_shape_infer_func)
+        onnx.defs.register_schema(schema)
+
+        # shape inference with registered schema
+        self.shape_infer_once(graph, N, La, Lb, out_len)
+
+        # clean up
+        onnx.defs.deregister_schema(schema.name, schema.since_version, schema.domain)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)