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

onnx/test/shape_inference_test.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 
 import itertools
 import unittest
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pytest
@@ -45,6 +45,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()
@@ -188,13 +191,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
@@ -976,6 +979,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)],
+        )  # type: ignore
+
     @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")
@@ -1680,6 +1699,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 +1812,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))]
+        )  # type: ignore
 
     def test_slice_with_input_shape_steps(self) -> None:
         graph = self._make_graph(
@@ -1900,7 +1938,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))]
+        )  # type: ignore
 
     def test_slice_negative_start(self) -> None:
         graph = self._make_graph(
@@ -1920,7 +1960,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))]
+        )  # type: ignore
 
     def test_slice_negative_step(self) -> None:
         graph = self._make_graph(
@@ -1942,7 +1984,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))]
+        )  # type: ignore
 
     def test_slice_variable_copy(self) -> None:
         graph = self._make_graph(
@@ -1960,7 +2004,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))]
+        )  # type: ignore
 
     def test_slice_variable_input_types(self) -> None:
         graph = self._make_graph(
@@ -2139,7 +2185,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))],
+        )  # type: ignore
 
     def test_conv_partial_missing_weight_shape(self) -> None:
         graph = self._make_graph(
@@ -2154,6 +2203,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"),
+                )
+            ],
+        )  # type: ignore
+
     def test_average_pool_auto_pads(self) -> None:
         graph = self._make_graph(
             [("x", TensorProto.FLOAT, (30, 4, 7, 6, 4))],
@@ -2231,7 +2319,9 @@ 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)
+                ),  # type: ignore
                 make_tensor_sequence_value_info(
                     "output_sequence", TensorProto.FLOAT, (2, None, 4)
                 ),
@@ -2441,7 +2531,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))]
+        )  # type: ignore
 
     def test_bernoulli_with_dtype(self) -> None:
         graph = self._make_graph(
@@ -2456,7 +2548,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))]
+        )  # type: ignore
 
     def _logical_binary_op(self, op: str, input_type: TensorProto.DataType) -> None:
         graph = self._make_graph(
@@ -2560,7 +2654,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))]
+        )  # type: ignore
 
     def test_space_to_depth(self) -> None:
         b = 10
@@ -2580,7 +2676,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))]
+        )  # type: ignore
 
     def test_depth_to_space(self) -> None:
         b = 10
@@ -2806,7 +2904,9 @@ 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)
+                ),  # type: ignore
                 make_tensor_value_info(
                     "z", TensorProto.INT64, (None, None, None, None)
                 ),
@@ -4409,7 +4509,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,))]
+        )  # type: ignore
 
     def test_if_no_shape_in_then_branch(self) -> None:
         then_graph = parse_graph(
@@ -4435,7 +4537,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)]
+        )  # type: ignore
 
     def test_if_no_shape_in_else_branch(self) -> None:
         then_graph = parse_graph(
@@ -4461,7 +4565,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)]
+        )  # type: ignore
 
     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
@@ -4580,7 +4686,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))]
+        )  # type: ignore
 
     def test_onehot_with_axis(self) -> None:
         graph = self._make_graph(
@@ -4592,7 +4700,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))]
+        )  # type: ignore
 
     def test_onehot_without_axis_2(self) -> None:
         graph = self._make_graph(
@@ -4605,7 +4715,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))]
+        )  # type: ignore
 
     def test_onehot_with_axis_2(self) -> None:
         graph = self._make_graph(
@@ -4618,7 +4730,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))]
+        )  # type: ignore
 
     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.
@@ -4913,7 +5027,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))],
+        )  # type: ignore
 
     def test_convineteger_partial_missing_weight_shape(self) -> None:
         graph = self._make_graph(
@@ -5136,7 +5253,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))],
+        )  # type: ignore
 
     def test_qlinearconv_partial_missing_weight_shape(self) -> None:
         graph = self._make_graph(
@@ -5567,7 +5687,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))]
+        )  # type: ignore
 
     @unittest.skipUnless(ONNX_ML, "ONNX_ML required to test ai.onnx.ml operators")
     def test_linearclassifier_1D_input(self) -> None:
@@ -5645,7 +5767,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))],
+        )  # type: ignore
 
     def test_roialign_symbolic_defaults(self) -> None:
         graph = self._make_graph(
@@ -5657,7 +5782,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))],
+        )  # type: ignore
 
     def test_roialign_num_rois(self) -> None:
         graph = self._make_graph(
@@ -5669,7 +5797,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))]
+        )  # type: ignore
+
+    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")
+                )
+            ],
+        )  # type: ignore
+
+    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")
+                )
+            ],
+        )  # type: ignore
 
     @parameterized.expand(
         all_versions_for("LabelEncoder") if ONNX_ML else [], skip_on_empty=True
@@ -5953,7 +6136,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)]
+        )  # type: ignore
 
     def test_constant_value_int(self) -> None:
         graph = self._make_graph(
@@ -6040,7 +6225,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,))]
+        )  # type: ignore
 
     def test_gathernd(self) -> None:
         graph = self._make_graph(
@@ -6084,11 +6271,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))]
+        )  # type: ignore
 
     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)]
+        )  # type: ignore
 
     def test_sequence_construct(self) -> None:
         graph = self._make_graph(
@@ -6261,7 +6452,9 @@ 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)
+                ),  # type: ignore
                 make_tensor_sequence_value_info(
                     "output_sequence", TensorProto.FLOAT, (2, None, None)
                 ),
@@ -6337,7 +6530,9 @@ 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)
+                ),  # type: ignore
                 make_tensor_value_info("output", TensorProto.FLOAT, (2, 3, None)),
             ],
         )  # type: ignore
@@ -6389,7 +6584,9 @@ 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")
+                ),  # type: ignore
                 make_tensor_sequence_value_info(
                     "output_sequence", TensorProto.FLOAT, (2, None, "x")
                 ),
@@ -6634,7 +6831,9 @@ 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")
+                ),  # type: ignore
                 make_tensor_value_info("out", TensorProto.FLOAT, (None, None, "x")),
             ],
         )  # type: ignore
@@ -6657,7 +6856,9 @@ 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")
+                ),  # type: ignore
                 make_tensor_value_info("out", TensorProto.FLOAT, (2, None, None)),
             ],
         )  # type: ignore
@@ -6680,7 +6881,9 @@ 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")
+                ),  # type: ignore
                 make_tensor_value_info("out", TensorProto.FLOAT, (None, None, "x")),
             ],
         )  # type: ignore
@@ -7023,7 +7226,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))]
+        )  # type: ignore
 
     def test_gatherelements_basic(self) -> None:
         graph = self._make_graph(
@@ -7044,7 +7249,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [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)]
+        )  # type: ignore
 
     def test_einsum_transpose(self) -> None:
         graph = self._make_graph(
@@ -7052,7 +7259,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))]
+        )  # type: ignore
 
     def test_einsum_dot(self) -> None:
         graph = self._make_graph(
@@ -7060,7 +7269,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, ())]
+        )  # type: ignore
 
     def test_einsum_scalar(self) -> None:
         graph = self._make_graph(
@@ -7068,7 +7279,9 @@ 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, ())]
+        )  # type: ignore
 
     def test_einsum_outer_prod(self) -> None:
         graph = self._make_graph(
@@ -7076,7 +7289,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))]
+        )  # type: ignore
 
     def test_einsum_sum_along_dim(self) -> None:
         graph = self._make_graph(
@@ -7084,7 +7299,9 @@ 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,))]
+        )  # type: ignore
 
     def test_einsum_ellipsis(self) -> None:
         graph = self._make_graph(
@@ -7092,7 +7309,9 @@ class TestShapeInference(TestShapeInferenceHelper):
             [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))]
+        )  # type: ignore
 
     def test_einsum_ellipsis_2(self) -> None:
         graph = self._make_graph(
@@ -7154,7 +7373,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))]
+        )  # type: ignore
 
     def test_einsum_left_hand_eqn(self) -> None:
         graph = self._make_graph(
@@ -7162,7 +7383,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))]
+        )  # type: ignore
 
     def test_einsum_incorrect_num_inputs(self) -> None:
         graph = self._make_graph(
@@ -7182,7 +7405,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,))]
+        )  # type: ignore
 
     def test_einsum_sum_A1(self) -> None:  # sums the values of A1
         graph = self._make_graph(
@@ -7190,7 +7415,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, ())]
+        )  # type: ignore
 
     def test_einsum_element_wise_multiplication_A1_B1(
         self,
@@ -7200,7 +7427,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,))]
+        )  # type: ignore
 
     def test_einsum_inner_product_A1_B1(self) -> None:  # inner product of A1 and B1
         graph = self._make_graph(
@@ -7208,7 +7437,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, ())]
+        )  # type: ignore
 
     def test_einsum_outer_product_A1_B1(self) -> None:  # outer product of A1 and B1
         graph = self._make_graph(
@@ -7216,7 +7447,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))]
+        )  # type: ignore
 
     def test_einsum_view_A2(self) -> None:  # returns a view of A2
         graph = self._make_graph(
@@ -7224,7 +7457,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))]
+        )  # type: ignore
 
     def test_einsum_view_A2_2(self) -> None:  # returns a view of A2, another case
         graph = self._make_graph(
@@ -7232,7 +7467,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))]
+        )  # type: ignore
 
     def test_einsum_transpose_A2(self) -> None:  # view transpose of A2
         graph = self._make_graph(
@@ -7240,7 +7477,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))]
+        )  # type: ignore
 
     def test_einsum_transpose_A2_to_ij(self) -> None:  # view transpose of A2
         graph = self._make_graph(
@@ -7248,7 +7487,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))]
+        )  # type: ignore
 
     def test_einsum_diag_A2(self) -> None:  # view main diagonal of A2
         graph = self._make_graph(
@@ -7256,7 +7497,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,))]
+        )  # type: ignore
 
     def test_einsum_trace_A2(self) -> None:  # sums main diagonal of A2
         graph = self._make_graph(
@@ -7264,7 +7507,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, ())]
+        )  # type: ignore
 
     def test_einsum_sum_A2(self) -> None:  # sums the values of A2
         graph = self._make_graph(
@@ -7272,7 +7517,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, ())]
+        )  # type: ignore
 
     def test_einsum_sum_columns_A2(
         self,
@@ -7282,7 +7529,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,))]
+        )  # type: ignore
 
     def test_einsum_sum_rows_A2(self) -> None:  # sum horizontally along the rows of A2
         graph = self._make_graph(
@@ -7290,7 +7539,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,))]
+        )  # type: ignore
 
     def test_einsum_element_wise_multiplication_A2_B2(
         self,
@@ -7300,7 +7551,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))]
+        )  # type: ignore
 
     def test_einsum_element_wise_multiplication_A2_B2_transpose(
         self,
@@ -7310,7 +7563,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))]
+        )  # type: ignore
 
     def test_einsum_matrix_multiplication_A2_B2(
         self,
@@ -7320,7 +7575,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))]
+        )  # type: ignore
 
     def test_einsum_matrix_multiplication_A2_B2_to_ik(
         self,
@@ -7330,7 +7587,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))]
+        )  # type: ignore
 
     def test_einsum_matrix_multiplication_A3_B3(
         self,
@@ -7340,7 +7599,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))]
+        )  # type: ignore
 
     def test_einsum_matrix_multiplication_A3_B3_transpose(
         self,
@@ -7350,7 +7611,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))]
+        )  # type: ignore
 
     def test_einsum_inner_product_A2_B2(self) -> None:  # inner product of A2 and B2
         graph = self._make_graph(
@@ -7358,7 +7621,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))]
+        )  # type: ignore
 
     def test_einsum_row_multiplication_A2_B2(
         self,
@@ -7368,7 +7633,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))]
+        )  # type: ignore
 
     def test_einsum_value_multiplication_A2_B2(
         self,
@@ -7378,7 +7645,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))]
+        )  # type: ignore
 
     def test_einsum_scalar_times_array(self) -> None:  # Scalar times array
         graph = self._make_graph(
@@ -7386,7 +7655,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))]
+        )  # type: ignore
 
     def test_einsum_matrix_vector_A2_B1(self) -> None:  # Matrix and vector.
         graph = self._make_graph(
@@ -7394,7 +7665,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,))]
+        )  # type: ignore
 
     def test_einsum_diag_multiplication_A2_B2(
         self,
@@ -7404,7 +7677,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,))]
+        )  # type: ignore
 
     def test_einsum_diag_dot_product_A2_B2(self) -> None:  # dot product of diagonals
         graph = self._make_graph(
@@ -7412,7 +7687,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, ())]
+        )  # type: ignore
 
     def test_negative_log_likehood_shape_is_NCdd(self) -> None:
         N, C = 3, 4
@@ -7428,7 +7705,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,))]
+        )  # type: ignore
 
     def test_negative_log_likehood_shape_is_NC_with_weight(self) -> None:
         N, C = 3, 4
@@ -7448,7 +7727,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,))]
+        )  # type: ignore
 
     def test_negative_log_likehood_shape_is_NC_reduction_mean(self) -> None:
         N, C = 3, 4
@@ -7464,7 +7745,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, ())]
+        )  # type: ignore
 
     def test_negative_log_likehood_shape_is_NC_with_weight_reduction_mean(self) -> None:
         N, C = 3, 4
@@ -7484,7 +7767,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, ())]
+        )  # type: ignore
 
     def test_negative_log_likehood_shape_is_NCd1d2(self) -> None:
         N, C, d1, d2 = 3, 4, 5, 6
@@ -7503,7 +7788,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))]
+        )  # type: ignore
 
     def test_negative_log_likehood_shape_is_NCd1d2_with_weight(self) -> None:
         N, C, d1, d2 = 3, 4, 5, 6
@@ -7523,7 +7810,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))]
+        )  # type: ignore
 
     def test_negative_log_likehood_shape_is_NCd1d2_reduction_sum(self) -> None:
         N, C, d1, d2 = 3, 4, 5, 6
@@ -7542,7 +7831,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, ())]
+        )  # type: ignore
 
     def test_negative_log_likehood_shape_is_NCd1d2_with_weight_reduction_mean(
         self,
@@ -7564,7 +7855,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, ())]
+        )  # type: ignore
 
     def test_negative_log_likehood_input_target_shape_mismatch(self) -> None:
         N, C, d1, d2 = 3, 4, 5, 6
@@ -7614,7 +7907,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,))]
+        )  # type: ignore
 
     def test_softmax_cross_entropy_mean(self) -> None:
         graph = self._make_graph(
@@ -7622,7 +7917,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, ())]
+        )  # type: ignore
 
     def test_softmax_cross_entropy_none_NCD1D2(self) -> None:
         graph = self._make_graph(
@@ -7633,7 +7930,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))]
+        )  # type: ignore
 
     def test_softmax_cross_entropy_mean_NCD1D2(self) -> None:
         graph = self._make_graph(
@@ -7644,7 +7943,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, ())]
+        )  # type: ignore
 
     def test_celu_function_output_shape(self) -> None:
         graph = self._make_graph(
@@ -7793,7 +8094,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))]
+        )  # type: ignore
 
     def test_trilu_lower(self) -> None:
         graph = self._make_graph(
@@ -7802,7 +8105,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))]
+        )  # type: ignore
 
     def test_trilu_upper_zero(self) -> None:
         graph = self._make_graph(
@@ -7811,7 +8116,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))]
+        )  # type: ignore
 
     def test_trilu_lower_one(self) -> None:
         graph = self._make_graph(
@@ -7819,7 +8126,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))]
+        )  # type: ignore
 
     def test_batch_norm_train(self) -> None:
         graph = self._make_graph(
@@ -7924,7 +8233,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))]
+        )  # type: ignore
 
     def test_batch_norm_test_no_dim(self) -> None:
         graph = self._make_graph(
@@ -7945,7 +8256,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)
+                )
+            ],
+        )  # type: ignore
 
     def test_batch_norm_train_no_shape(self) -> None:
         graph = self._make_graph(
@@ -7981,13 +8299,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))]
+        )  # type: ignore
 
     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))]
+        )  # type: ignore
 
     def test_nonzero_existing_dim_param(self) -> None:
         graph = self._make_graph(
@@ -7995,13 +8317,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"))]
+        )  # type: ignore
 
     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))]
+        )  # type: ignore
 
     def test_optional_construct_empty_tensor(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(
@@ -8280,7 +8606,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]
+        )  # type: ignore
 
     def test_where_bfloat(self) -> None:
         graph = self._make_graph(
@@ -8292,7 +8620,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,))]
+        )  # type: ignore
 
     def test_parse_data_with_unsupported_tensor_type(self) -> None:
         model = helper.make_model(
@@ -9111,7 +9441,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,  # type: ignore[arg-type]
+                    ),
                 ]
                 value_infos = [
                     make_tensor_value_info("dft_length", TensorProto.INT64, ()),
@@ -9127,7 +9462,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,  # type: ignore[arg-type]
+                    ),
                 ]
                 value_infos = [
                     make_tensor_value_info("dft_length", TensorProto.INT64, ())
@@ -9795,6 +10135,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 +10165,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 +10279,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 +10338,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:
@@ -10083,6 +10452,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 +10535,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,))]
+        )  # type: ignore
 
     def test_compress_with_axis(self) -> None:
         graph = self._make_graph(
@@ -10143,7 +10548,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))],
+        )  # type: ignore
 
     def test_check_type_when_schema_has_empty_io(self):
         input = """
@@ -10232,6 +10640,28 @@ 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)],
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)