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

onnx/test/shape_inference_test.py

@@ -53,7 +53,7 @@ def get_available_versions(schema: OpSchema) -> set[int]:
             versions.add(
                 defs.get_schema(schema.name, version, schema.domain).since_version
             )
-        except SchemaError:
+        except SchemaError:  # noqa: PERF203
             break
     return versions
 
@@ -399,7 +399,7 @@ class TestShapeInference(TestShapeInferenceHelper):
     @unittest.skip(
         "Issue #5960"
     )  # FIXME(#5960) propagateElemTypeFromAttributeToOutput does not validate against output type constraints
-    def test_cast_to_complex(self, _, version) -> None:
+    def test_cast_to_complex(self, _, version) -> None:  # noqa: ARG002
         graph = self._make_graph(
             [("x", TensorProto.FLOAT, (2, 4, 3))],
             [make_node("Cast", ["x"], ["y"], to=TensorProto.COMPLEX128)],
@@ -4594,6 +4594,32 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         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(
+            [
+                ("indices", TensorProto.INT64, (2, 2)),
+                ("depth", TensorProto.INT64, ()),
+                ("values", TensorProto.FLOAT, (2,)),
+            ],
+            [make_node("OneHot", ["indices", "depth", "values"], "Y")],
+            [],
+            initializer=[make_tensor("depth", TensorProto.INT64, (), (256,))],
+        )
+        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(
+            [
+                ("indices", TensorProto.INT64, (2, 3, 5)),
+                ("depth", TensorProto.INT64, (1,)),
+                ("values", TensorProto.FLOAT, (2,)),
+            ],
+            [make_node("OneHot", ["indices", "depth", "values"], "Y", axis=1)],
+            [],
+            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
+
     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.
         # this breaks the subgraph inferencing as it expects the number of inputs passed from Loop to match
@@ -7026,7 +7052,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="ij->ji")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (None, None))])  # 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(
@@ -7050,7 +7076,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,ab->ijab")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (None, None, None, None))])  # 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(
@@ -7058,7 +7084,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="i j->i ")],
             [],
         )
-        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, (3,))])  # type: ignore
 
     def test_einsum_ellipsis(self) -> None:
         graph = self._make_graph(
@@ -7066,26 +7092,36 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x"], ["y"], equation="... ii ->... i")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (None, None))])  # 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(
-            [("x", TensorProto.FLOAT, (2, 2, 2)), ("y", TensorProto.FLOAT, (2, 2, 2))],
+            [("x", TensorProto.FLOAT, (2, 3, 4)), ("y", TensorProto.FLOAT, (2, 4, 5))],
             [make_node("Einsum", ["x", "y"], ["z"], equation="...ij,...jk->...ik")],
             [],
         )
         self._assert_inferred(
-            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (None, None, None))]
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 5))]
         )  # type: ignore
 
     def test_einsum_ellipsis_3(self) -> None:
         graph = self._make_graph(
-            [("x", TensorProto.FLOAT, (2, 2, 2)), ("y", TensorProto.FLOAT, (2, 2, 2))],
+            [("x", TensorProto.FLOAT, (2, 3, 4)), ("y", TensorProto.FLOAT, (2, 4, 5))],
             [make_node("Einsum", ["x", "y"], ["z"], equation="...ij,...jk")],
             [],
         )
         self._assert_inferred(
-            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (None, None, None))]
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (2, 3, 5))]
+        )  # type: ignore
+
+    def test_einsum_ellipsis_broadcast(self) -> None:
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (1, 3, 4)), ("y", TensorProto.FLOAT, (32, 4, 5))],
+            [make_node("Einsum", ["x", "y"], ["z"], equation="...ij,...jk->...ik")],
+            [],
+        )
+        self._assert_inferred(
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (32, 3, 5))]
         )  # type: ignore
 
     def test_einsum_contraction(self) -> None:
@@ -7099,11 +7135,7 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self._assert_inferred(
             graph,
-            [
-                make_tensor_value_info(
-                    "z", TensorProto.FLOAT, (None, None, None, None, None)
-                )
-            ],
+            [make_tensor_value_info("z", TensorProto.FLOAT, (5, 6, 7, 9, 10))],
         )  # type: ignore
 
     def test_einsum_contraction_2(self) -> None:
@@ -7113,7 +7145,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             [],
         )
         self._assert_inferred(
-            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (None, None))]
+            graph, [make_tensor_value_info("z", TensorProto.FLOAT, (4, 5))]
         )  # type: ignore
 
     def test_einsum_batch_matmul(self) -> None:
@@ -7122,7 +7154,7 @@ 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, (None, None, None))])  # 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(
@@ -7130,7 +7162,7 @@ class TestShapeInference(TestShapeInferenceHelper):
             [make_node("Einsum", ["x", "y"], ["z"], equation="ij,kl")],
             [],
         )
-        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (None, None, None, None))])  # 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(
@@ -7144,6 +7176,244 @@ class TestShapeInference(TestShapeInferenceHelper):
         )
         self.assertRaises(onnx.shape_inference.InferenceError, self._inferred, graph)
 
+    def test_einsum_view_A1(self) -> None:  # returns a view of A1
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3,))],
+            [make_node("Einsum", ["x"], ["y"], equation="i")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3,))],
+            [make_node("Einsum", ["x"], ["y"], equation="i->")],
+            [],
+        )
+        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, ())])  # type: ignore
+
+    def test_einsum_element_wise_multiplication_A1_B1(
+        self,
+    ) -> None:  # element-wise multiplication of A1 and B1
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3,)), ("y", TensorProto.FLOAT, (3,))],
+            [make_node("Einsum", ["x", "y"], ["z"], equation="i,i->i")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3,)), ("y", TensorProto.FLOAT, (3,))],
+            [make_node("Einsum", ["x", "y"], ["z"], equation="i,i->")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3,)), ("y", TensorProto.FLOAT, (3,))],
+            [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
+
+    def test_einsum_view_A2(self) -> None:  # returns a view of A2
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x"], ["y"], equation="ij->ij")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x"], ["y"], equation="ij")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x"], ["y"], equation="ji")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x"], ["y"], equation="ji->ij")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x"], ["y"], equation="ii->i")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x"], ["y"], equation="ii->")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x"], ["y"], equation="ij->")],
+            [],
+        )
+        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, ())])  # type: ignore
+
+    def test_einsum_sum_columns_A2(
+        self,
+    ) -> None:  # sum down the columns of A2 (across rows)
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x"], ["y"], equation="ij->j")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x"], ["y"], equation="ij->i")],
+            [],
+        )
+        self._assert_inferred(graph, [make_tensor_value_info("y", TensorProto.FLOAT, (3,))])  # type: ignore
+
+    def test_einsum_element_wise_multiplication_A2_B2(
+        self,
+    ) -> None:  # element-wise multiplication of A2 and B2
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3, 3))],
+            [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
+
+    def test_einsum_element_wise_multiplication_A2_B2_transpose(
+        self,
+    ) -> None:  # element-wise multiplication of A2 and B2.T
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3, 3))],
+            [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
+
+    def test_einsum_matrix_multiplication_A2_B2(
+        self,
+    ) -> None:  # matrix multiplication of A2 and B2
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x", "y"], ["z"], equation="ij,jk")],
+            [],
+        )
+        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,
+    ) -> None:  # matrix multiplication of A2 and B2
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3, 3))],
+            [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
+
+    def test_einsum_matrix_multiplication_A3_B3(
+        self,
+    ) -> None:  # matrix multiplication of A3 and B3 (a stack of 2D matrices)
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (2, 3, 3)), ("y", TensorProto.FLOAT, (2, 3, 3))],
+            [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
+
+    def test_einsum_matrix_multiplication_A3_B3_transpose(
+        self,
+    ) -> None:  # matrix multiplication of A3 and B3 (a stack of 2D matrices)
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (2, 3, 3)), ("y", TensorProto.FLOAT, (2, 3, 3))],
+            [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
+
+    def test_einsum_inner_product_A2_B2(self) -> None:  # inner product of A2 and B2
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3, 3))],
+            [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
+
+    def test_einsum_row_multiplication_A2_B2(
+        self,
+    ) -> None:  # each row of A2 multiplied by B2
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3, 3))],
+            [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
+
+    def test_einsum_value_multiplication_A2_B2(
+        self,
+    ) -> None:  # each value of A2 multiplied by B2
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3, 3))],
+            [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
+
+    def test_einsum_scalar_times_array(self) -> None:  # Scalar times array
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, ()), ("y", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x", "y"], ["z"], equation=",ij->ij")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3,))],
+            [make_node("Einsum", ["x", "y"], ["z"], equation="ij,j->i")],
+            [],
+        )
+        self._assert_inferred(graph, [make_tensor_value_info("z", TensorProto.FLOAT, (3,))])  # type: ignore
+
+    def test_einsum_diag_multiplication_A2_B2(
+        self,
+    ) -> None:  # diagonals multiplied by each other
+        graph = self._make_graph(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x", "y"], ["z"], equation="ii,ii->i")],
+            [],
+        )
+        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(
+            [("x", TensorProto.FLOAT, (3, 3)), ("y", TensorProto.FLOAT, (3, 3))],
+            [make_node("Einsum", ["x", "y"], ["z"], equation="ii,ii->")],
+            [],
+        )
+        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
         graph = self._make_graph(
@@ -9902,6 +10172,66 @@ class TestShapeInference(TestShapeInferenceHelper):
             op_schema.name, op_schema.since_version, op_schema.domain
         )
 
+    def test_issue_layer_normalization_6187(self):
+        modeltxt = """
+        <
+        ir_version: 10,
+        opset_import: ["" : 17]
+        >
+        graph (float in0, float[2,7,8,1,3] in1, float[3,7] in2) => () {
+        out0, out1, out2 = LayerNormalization <epsilon: float = -841.058, stash_type: int = -940> (in0, in1, in2)
+        }
+        """
+        model = onnx.parser.parse_model(modeltxt)
+        with self.assertRaises(onnx.shape_inference.InferenceError):
+            onnx.checker.check_model(model, full_check=True)
+            onnx.shape_inference.infer_shapes(model)
+
+    def test_issue_conv_6180(self):
+        modeltxt = """
+        <
+        ir_version: 9,
+        opset_import: ["" : 11]
+        >
+        graph (float[7,6,1,5] in0, float in1, float[7,2,3,2,1] in2) => () {
+        out0 = Conv <auto_pad = "NOTSET", group = 1> (in0, in1, in2)
+        }
+        """
+        model = onnx.parser.parse_model(modeltxt)
+        with self.assertRaises(onnx.shape_inference.InferenceError):
+            onnx.checker.check_model(model, full_check=True)
+            onnx.shape_inference.infer_shapes(model)
+
+    def test_issue_gemm_6185(self):
+        modeltxt = """
+        <
+        ir_version: 10,
+        opset_import: ["" : 6]
+        >
+        graph (double[2,1] in0, double in1, double[2] in2) => () {
+        out0 = Gemm <alpha: float = 1, beta: float = -693.752, broadcast: int = -436, transB: int = 823> (in0, in1, in2)
+        }
+        """
+        model = onnx.parser.parse_model(modeltxt)
+        with self.assertRaises(onnx.checker.ValidationError):
+            onnx.checker.check_model(model, full_check=True)
+            onnx.shape_inference.infer_shapes(model)
+
+    def test_issue_stft_6186(self):
+        modeltxt = """
+        <
+        ir_version: 10,
+        opset_import: ["" : 17]
+        >
+        graph (float16[3] in0, int32[2] in1, float16[7,8,8,8] in2, int32[8,1,7,2] in3) => () {
+        out0 = STFT (in0, in1, in2, in3)
+        }
+        """
+        model = onnx.parser.parse_model(modeltxt)
+        with self.assertRaises(onnx.shape_inference.InferenceError):
+            onnx.checker.check_model(model, full_check=True)
+            onnx.shape_inference.infer_shapes(model)
+
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main(verbosity=2)

onnx/test/symbolic_shape_test.py

@@ -1,9 +1,9 @@
 # Copyright (c) ONNX Project Contributors
 
 # SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
 
 import unittest
-from typing import List, Optional
 
 import onnx.shape_inference
 from onnx import ModelProto, TensorProto, TensorShapeProto, ValueInfoProto, helper
@@ -12,7 +12,7 @@ from onnx.helper import make_model, make_tensor_value_info
 
 class TestSymbolicShape(unittest.TestCase):
     def _assert_valueinfo_shape(
-        self, onnx_model: ModelProto, value_infos: List[ValueInfoProto]
+        self, onnx_model: ModelProto, value_infos: list[ValueInfoProto]
     ) -> None:
         """Assert onnx_model.value_info should be the same as expected value_infos
         Instead of exact symbol, use -1 to represent symbolic shape in expected value_infos
@@ -48,7 +48,7 @@ class TestSymbolicShape(unittest.TestCase):
 
     def _get_shape_from_name(
         self, onnx_model: ModelProto, name: str
-    ) -> Optional[TensorShapeProto]:
+    ) -> TensorShapeProto | None:
         """Get shape from tensor_type or sparse_tensor_type according to given name"""
         inputs = list(onnx_model.graph.input)
         outputs = list(onnx_model.graph.output)