onnx 1.19.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl → 1.19.1rc1__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl

onnx/__init__.py

@@ -143,6 +143,9 @@ from onnx import (
     version_converter,
 )
 
+if typing.TYPE_CHECKING:
+    from collections.abc import Sequence
+
 __version__ = onnx.version.version
 
 # Supported model formats that can be loaded from and saved to
@@ -368,3 +371,98 @@ def save_tensor(
 load = load_model
 load_from_string = load_model_from_string
 save = save_model
+
+
+def _model_proto_repr(self: ModelProto) -> str:
+    if self.domain:
+        domain = f", domain='{self.domain}'"
+    else:
+        domain = ""
+    if self.producer_name:
+        producer_name = f", producer_name='{self.producer_name}'"
+    else:
+        producer_name = ""
+    if self.producer_version:
+        producer_version = f", producer_version='{self.producer_version}'"
+    else:
+        producer_version = ""
+    if self.graph:
+        graph = f", graph={self.graph!r}"
+    else:
+        graph = ""
+    if self.functions:
+        functions = f", functions=<{len(self.functions)} functions>"
+    else:
+        functions = ""
+    if self.opset_import:
+        opset_import = f", opset_import={_operator_set_protos_repr(self.opset_import)}"
+    else:
+        opset_import = ""
+    return f"ModelProto(ir_version={self.ir_version}{opset_import}{domain}{producer_name}{producer_version}{graph}{functions})"
+
+
+def _graph_proto_repr(self: GraphProto) -> str:
+    if self.initializer:
+        initializer = f", initializer=<{len(self.initializer)} initializers>"
+    else:
+        initializer = ""
+    if self.node:
+        node = f", node=<{len(self.node)} nodes>"
+    else:
+        node = ""
+    if self.value_info:
+        value_info = f", value_info=<{len(self.value_info)} value_info>"
+    else:
+        value_info = ""
+    if self.input:
+        input = f", input=<{len(self.input)} inputs>"
+    else:
+        input = ""
+    if self.output:
+        output = f", output=<{len(self.output)} outputs>"
+    else:
+        output = ""
+    return f"GraphProto('{self.name}'{input}{output}{initializer}{node}{value_info})"
+
+
+def _function_proto_repr(self: FunctionProto) -> str:
+    if self.domain:
+        domain = f", domain='{self.domain}'"
+    else:
+        domain = ""
+    if self.overload:
+        overload = f", overload='{self.overload}'"
+    else:
+        overload = ""
+    if self.node:
+        node = f", node=<{len(self.node)} nodes>"
+    else:
+        node = ""
+    if self.attribute:
+        attribute = f", attribute={self.attribute}"
+    else:
+        attribute = ""
+    if self.opset_import:
+        opset_import = f", opset_import={_operator_set_protos_repr(self.opset_import)}"
+    else:
+        opset_import = ""
+    if self.input:
+        input = f", input=<{len(self.input)} inputs>"
+    else:
+        input = ""
+    if self.output:
+        output = f", output=<{len(self.output)} outputs>"
+    else:
+        output = ""
+    return f"FunctionProto('{self.name}'{domain}{overload}{opset_import}{input}{output}{attribute}{node})"
+
+
+def _operator_set_protos_repr(protos: Sequence[OperatorSetIdProto]) -> str:
+    opset_imports = {proto.domain: proto.version for proto in protos}
+    return repr(opset_imports)
+
+
+# Override __repr__ for some proto classes to make it more efficient
+ModelProto.__repr__ = _model_proto_repr  # type: ignore[method-assign,assignment]
+GraphProto.__repr__ = _graph_proto_repr  # type: ignore[method-assign,assignment]
+FunctionProto.__repr__ = _function_proto_repr  # type: ignore[method-assign,assignment]

onnx/backend/test/case/node/__init__.py

@@ -18,6 +18,7 @@ from onnx.onnx_pb import (
     GraphProto,
     ModelProto,
     NodeProto,
+    OperatorSetIdProto,
     TensorProto,
     TypeProto,
 )
@@ -128,11 +129,25 @@ def function_expand_helper(
 
 
 def function_testcase_helper(
-    node: NodeProto, input_types: list[TypeProto], name: str
+    node: NodeProto,
+    input_types: list[TypeProto],
+    name: str,
+    opset_imports: Sequence[OperatorSetIdProto] | None = None,
 ) -> tuple[list[tuple[list[NodeProto], Any]], int]:
     test_op = node.op_type
     op_prefix = test_op + "_" + name + "_expanded_function_"
-    schema = onnx.defs.get_schema(test_op, domain=node.domain)
+    if opset_imports is None:
+        # No opset in the model. We take the most recent definition.
+        schema = onnx.defs.get_schema(test_op, domain=node.domain)
+    else:
+        # We take the function coming defined in the specific version mentioned
+        # in the model.
+        if len(opset_imports) != 1:
+            raise ValueError(
+                f"Only one domain is allowed but {len(opset_imports)} found."
+            )
+        version = opset_imports[0].version
+        schema = onnx.defs.get_schema(test_op, version, domain=node.domain)
 
     # an op schema may have several functions, each for one opset version
     # opset versions include the op's since_version and other opset versions
@@ -327,7 +342,9 @@ def expect(
     (
         expanded_tests,
         since_version,
-    ) = function_testcase_helper(node, merged_types, name)
+    ) = function_testcase_helper(
+        node, merged_types, name, opset_imports=kwargs.get("opset_imports")
+    )
     for expanded_function_nodes, func_opset_import in expanded_tests:
         kwargs["producer_name"] = "backend-test"
 

onnx/backend/test/case/node/attention.py

@@ -27,6 +27,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -44,6 +45,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_fp16",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -61,6 +63,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_gqa",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -78,6 +81,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_diff_heads_sizes",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -101,6 +105,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_scaled",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -124,6 +129,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_gqa_scaled",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -147,6 +153,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_diff_heads_sizes_scaled",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -169,6 +176,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -191,6 +199,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_gqa_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -218,6 +227,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_diff_heads_sizes_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -245,6 +255,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_4d_attn_mask",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -272,6 +283,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_4d_attn_mask_3d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -301,6 +313,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_4d_attn_mask_3d_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -328,6 +341,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_4d_attn_mask_4d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -357,6 +371,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_4d_attn_mask_4d_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -384,6 +399,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_4d_attn_mask_bool",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -411,6 +427,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_4d_attn_mask_bool_4d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -438,6 +455,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_4d_gqa_attn_mask",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -465,6 +483,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_4d_diff_heads_sizes_attn_mask",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -497,6 +516,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value],
             name="test_attention_4d_with_past_and_present",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -529,6 +549,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value],
             name="test_attention_4d_gqa_with_past_and_present",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -561,6 +582,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value],
             name="test_attention_4d_gqa_with_past_and_present_fp16",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -593,6 +615,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value],
             name="test_attention_4d_diff_heads_with_past_and_present",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -625,6 +648,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value],
             name="test_attention_4d_diff_heads_with_past_and_present_mask3d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -657,6 +681,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value],
             name="test_attention_4d_diff_heads_with_past_and_present_mask4d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -679,6 +704,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_softcap",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -701,6 +727,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_gqa_softcap",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -728,6 +755,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d_diff_heads_sizes_softcap",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -749,6 +777,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y, qk_matmul_output],
             name="test_attention_4d_with_qk_matmul",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -778,6 +807,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y, qk_matmul_output],
             name="test_attention_4d_with_qk_matmul_bias",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -809,6 +839,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y, qk_matmul_output],
             name="test_attention_4d_with_qk_matmul_softcap",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -838,6 +869,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y, qk_matmul_output],
             name="test_attention_4d_with_qk_matmul_softmax",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -872,6 +904,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_4d_with_past_and_present_qk_matmul_bias",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -906,6 +939,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -940,6 +974,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -976,6 +1011,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_4d_with_past_and_present_qk_matmul_bias_3d_mask_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1012,6 +1048,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_4d_with_past_and_present_qk_matmul_bias_4d_mask_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1044,6 +1081,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_4d_with_past_and_present_qk_matmul",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1074,6 +1112,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1104,6 +1143,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_gqa",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1134,6 +1174,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_diff_heads_sizes",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1167,6 +1208,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_scaled",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1200,6 +1242,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_gqa_scaled",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1233,6 +1276,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_diff_heads_sizes_scaled",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1265,6 +1309,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1297,6 +1342,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_gqa_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1329,6 +1375,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_diff_heads_sizes_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1361,6 +1408,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_3d_attn_mask",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1393,6 +1441,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_3d_gqa_attn_mask",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1425,6 +1474,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask],
             outputs=[Y],
             name="test_attention_3d_diff_heads_sizes_attn_mask",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1457,6 +1507,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_softcap",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1489,6 +1540,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_gqa_softcap",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1521,6 +1573,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_diff_heads_sizes_softcap",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1558,6 +1611,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value],
             name="test_attention_3d_with_past_and_present",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1595,6 +1649,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value],
             name="test_attention_3d_gqa_with_past_and_present",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1632,6 +1687,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value],
             name="test_attention_3d_diff_heads_with_past_and_present",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1669,6 +1725,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_3d_with_past_and_present_qk_matmul",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1708,6 +1765,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_3d_with_past_and_present_qk_matmul_bias",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1749,6 +1807,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_3d_with_past_and_present_qk_matmul_softcap",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1788,6 +1847,7 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, past_key, past_value],
             outputs=[Y, present_key, present_value, qk_matmul_output],
             name="test_attention_3d_with_past_and_present_qk_matmul_softmax",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1842,6 +1902,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d_transpose_verification",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -1871,4 +1932,5 @@ class Attention(Base):
             inputs=[Q, K, V, attn_mask, nonpad_kv_seqlen],
             outputs=[Y],
             name="test_attention_4d_diff_heads_mask4d_padded_kv",
+            opset_imports=[onnx.helper.make_opsetid("", 24)],
         )

onnx/backend/test/case/node/rotaryembedding.py

@@ -106,8 +106,8 @@ class RotaryEmbedding(Base):
 
         input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
         position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
-        sin_cache_data = np.random.rand(50, 4).astype(np.float32)
-        cos_cache_data = np.random.rand(50, 4).astype(np.float32)
+        sin_cache_data = np.random.rand(50, 2).astype(np.float32)
+        cos_cache_data = np.random.rand(50, 2).astype(np.float32)
 
         expected_output = rotary_embedding(
             input_data,
@@ -136,8 +136,8 @@ class RotaryEmbedding(Base):
 
         input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
         position_ids_data = np.random.uniform(0, 50, (2, 3)).astype(np.int64)
-        sin_cache_data = np.random.rand(50, 4).astype(np.float32)
-        cos_cache_data = np.random.rand(50, 4).astype(np.float32)
+        sin_cache_data = np.random.rand(50, 2).astype(np.float32)
+        cos_cache_data = np.random.rand(50, 2).astype(np.float32)
 
         expected_output = rotary_embedding(
             input_data,
@@ -213,8 +213,8 @@ class RotaryEmbedding(Base):
         )
 
         input_data = np.random.rand(2, 4, 3, 8).astype(np.float32)
-        sin_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
-        cos_cache_data = np.random.rand(2, 3, 4).astype(np.float32)
+        sin_cache_data = np.random.rand(2, 3, 2).astype(np.float32)
+        cos_cache_data = np.random.rand(2, 3, 2).astype(np.float32)
 
         expected_output = rotary_embedding(
             input_data,