onnx 1.18.0__cp311-cp311-win_arm64.whl → 1.19.1rc1__cp311-cp311-win_arm64.whl

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

@@ -27,6 +27,25 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_4d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
+        )
+
+    @staticmethod
+    def export_attention_fp16() -> None:
+        node = onnx.helper.make_node("Attention", inputs=["Q", "K", "V"], outputs=["Y"])
+
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float16)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float16)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float16)
+
+        Y, _, _, _ = _compute_attention(Q, K, V)
+
+        expect(
+            node,
+            inputs=[Q, K, V],
+            outputs=[Y],
+            name="test_attention_4d_fp16",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -44,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
@@ -61,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
@@ -84,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
@@ -107,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
@@ -130,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
@@ -152,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
@@ -174,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
@@ -201,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
@@ -228,6 +255,123 @@ 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
+    def export_attention_attn_3d_mask() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask"],
+            outputs=["Y"],
+        )
+
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        attn_mask = np.random.rand(2, 1, 4, 6).astype(np.float32)
+
+        Y, _, _, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+        )
+
+        expect(
+            node,
+            inputs=[Q, K, V, attn_mask],
+            outputs=[Y],
+            name="test_attention_4d_attn_mask_3d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
+        )
+
+    @staticmethod
+    def export_attention_attn_3d_mask_causal() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask"],
+            outputs=["Y"],
+            is_causal=1,
+        )
+
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        attn_mask = np.random.rand(2, 1, 4, 6).astype(np.float32)
+
+        Y, _, _, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            is_causal=1,
+        )
+
+        expect(
+            node,
+            inputs=[Q, K, V, attn_mask],
+            outputs=[Y],
+            name="test_attention_4d_attn_mask_3d_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
+        )
+
+    @staticmethod
+    def export_attention_attn_4d_mask() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask"],
+            outputs=["Y"],
+        )
+
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        attn_mask = np.random.rand(2, 3, 4, 6).astype(np.float32)
+
+        Y, _, _, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+        )
+
+        expect(
+            node,
+            inputs=[Q, K, V, attn_mask],
+            outputs=[Y],
+            name="test_attention_4d_attn_mask_4d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
+        )
+
+    @staticmethod
+    def export_attention_attn_4d_mask_causal() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask"],
+            outputs=["Y"],
+            is_causal=1,
+        )
+
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        attn_mask = np.random.rand(2, 3, 4, 6).astype(np.float32)
+
+        Y, _, _, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            is_causal=1,
+        )
+
+        expect(
+            node,
+            inputs=[Q, K, V, attn_mask],
+            outputs=[Y],
+            name="test_attention_4d_attn_mask_4d_causal",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -241,7 +385,7 @@ class Attention(Base):
         Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
         K = np.random.rand(2, 3, 6, 8).astype(np.float32)
         V = np.random.rand(2, 3, 6, 8).astype(np.float32)
-        attn_mask = np.random.rand(4, 6).astype(np.bool)
+        attn_mask = np.random.rand(4, 6).astype(bool)
 
         Y, _, _, _ = _compute_attention(
             Q,
@@ -255,6 +399,35 @@ 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
+    def export_attention_attn_mask_bool_4d() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask"],
+            outputs=["Y"],
+        )
+
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        attn_mask = np.random.rand(2, 3, 4, 6).astype(bool)
+
+        Y, _, _, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+        )
+
+        expect(
+            node,
+            inputs=[Q, K, V, attn_mask],
+            outputs=[Y],
+            name="test_attention_4d_attn_mask_bool_4d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -282,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
@@ -309,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
@@ -341,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
@@ -373,6 +549,40 @@ 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
+    def export_attention_gqa_with_past_and_present_fp16() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
+            outputs=["Y", "present_key", "present_value"],
+        )
+
+        past_sequence_length = 12
+        Q = np.random.rand(2, 9, 4, 8).astype(np.float16)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float16)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float16)
+        attn_mask = np.random.rand(4, 6 + past_sequence_length).astype(np.float16)
+        past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float16)
+        past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float16)
+
+        Y, present_key, present_value, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            past_key=past_key,
+            past_value=past_value,
+        )
+
+        expect(
+            node,
+            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
@@ -405,6 +615,73 @@ 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
+    def export_attention_diff_head_sizes_with_past_and_present_mask3D() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
+            outputs=["Y", "present_key", "present_value"],
+        )
+
+        past_sequence_length = 12
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 10).astype(np.float32)
+        attn_mask = np.random.rand(2, 1, 4, 6 + past_sequence_length).astype(np.float32)
+        past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+        past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32)
+
+        Y, present_key, present_value, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            past_key=past_key,
+            past_value=past_value,
+        )
+
+        expect(
+            node,
+            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
+    def export_attention_diff_head_sizes_with_past_and_present_mask4D() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
+            outputs=["Y", "present_key", "present_value"],
+        )
+
+        past_sequence_length = 12
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 10).astype(np.float32)
+        attn_mask = np.random.rand(2, 3, 4, 6 + past_sequence_length).astype(np.float32)
+        past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+        past_value = np.random.rand(2, 3, past_sequence_length, 10).astype(np.float32)
+
+        Y, present_key, present_value, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            past_key=past_key,
+            past_value=past_value,
+        )
+
+        expect(
+            node,
+            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
@@ -427,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
@@ -449,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
@@ -476,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
@@ -497,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
@@ -526,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
@@ -557,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
@@ -586,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
@@ -620,6 +904,151 @@ 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
+    def export_attention_with_past_and_present_qk_matmul_bias_3d_mask() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
+            outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
+            qk_matmul_output_mode=1,
+        )
+
+        past_sequence_length = 12
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        attn_mask = np.random.rand(2, 1, 4, 6 + past_sequence_length).astype(np.float32)
+        past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+        past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+
+        Y, present_key, present_value, qk_matmul_output = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            past_key=past_key,
+            past_value=past_value,
+            qk_matmul_output_mode=1,
+        )
+
+        expect(
+            node,
+            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
+    def export_attention_with_past_and_present_qk_matmul_bias_4d_mask() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
+            outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
+            qk_matmul_output_mode=1,
+        )
+
+        past_sequence_length = 12
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        attn_mask = np.random.rand(2, 3, 4, 6 + past_sequence_length).astype(np.float32)
+        past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+        past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+
+        Y, present_key, present_value, qk_matmul_output = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            past_key=past_key,
+            past_value=past_value,
+            qk_matmul_output_mode=1,
+        )
+
+        expect(
+            node,
+            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
+    def export_attention_with_past_and_present_qk_matmul_bias_3d_mask_causal() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
+            outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
+            qk_matmul_output_mode=1,
+            is_causal=1,
+        )
+
+        past_sequence_length = 12
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        attn_mask = np.random.rand(2, 1, 4, 6 + past_sequence_length).astype(np.float32)
+        past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+        past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+
+        Y, present_key, present_value, qk_matmul_output = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            past_key=past_key,
+            past_value=past_value,
+            qk_matmul_output_mode=1,
+            is_causal=1,
+        )
+
+        expect(
+            node,
+            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
+    def export_attention_with_past_and_present_qk_matmul_bias_4d_mask_causal() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask", "past_key", "past_value"],
+            outputs=["Y", "present_key", "present_value", "qk_matmul_output"],
+            qk_matmul_output_mode=1,
+            is_causal=1,
+        )
+
+        past_sequence_length = 12
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        attn_mask = np.random.rand(2, 3, 4, 6 + past_sequence_length).astype(np.float32)
+        past_key = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+        past_value = np.random.rand(2, 3, past_sequence_length, 8).astype(np.float32)
+
+        Y, present_key, present_value, qk_matmul_output = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            past_key=past_key,
+            past_value=past_value,
+            qk_matmul_output_mode=1,
+            is_causal=1,
+        )
+
+        expect(
+            node,
+            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
@@ -652,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
@@ -682,6 +1112,7 @@ class Attention(Base):
             inputs=[Q, K, V],
             outputs=[Y],
             name="test_attention_3d",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
         )
 
     @staticmethod
@@ -712,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
@@ -742,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
@@ -775,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
@@ -808,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
@@ -841,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
@@ -873,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
@@ -905,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
@@ -937,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
@@ -969,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
@@ -1001,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
@@ -1033,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
@@ -1065,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
@@ -1097,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
@@ -1129,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
@@ -1166,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
@@ -1203,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
@@ -1240,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
@@ -1277,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
@@ -1316,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
@@ -1357,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
@@ -1396,4 +1847,90 @@ 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
+    def export_attention_3d_transpose_verification() -> None:
+        """Test case to verify correct 3D to 4D transpose behavior.
+
+        This test verifies that 3D inputs are correctly reshaped and transposed
+        according to the ONNX specification:
+        [batch_size, seq_length, hidden_size] ->
+        [batch_size, seq_length, num_heads, head_size] ->
+        [batch_size, num_heads, seq_length, head_size]
+        """
+        q_num_heads, kv_num_heads = 3, 3
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V"],
+            outputs=["Y"],
+            q_num_heads=q_num_heads,
+            kv_num_heads=kv_num_heads,
+        )
+
+        # Test inputs that will clearly demonstrate the transpose behavior
+        batch_size = 1
+        q_seq_length = 2
+        kv_seq_length = 2
+        head_size = 4
+        q_hidden_size = q_num_heads * head_size  # 3 * 4 = 12
+        kv_hidden_size = kv_num_heads * head_size  # 3 * 4 = 12
+
+        # Create structured inputs to verify correct transpose behavior
+        # Q has a pattern where each position in hidden dimension has a specific value
+        Q = np.zeros((batch_size, q_seq_length, q_hidden_size), dtype=np.float32)
+        # Fill Q with pattern: head0=[1,1,1,1], head1=[2,2,2,2], head2=[3,3,3,3]
+        for head in range(q_num_heads):
+            start_idx = head * head_size
+            end_idx = start_idx + head_size
+            Q[0, :, start_idx:end_idx] = float(head + 1)
+
+        K = np.ones((batch_size, kv_seq_length, kv_hidden_size), dtype=np.float32) * 0.1
+        V = np.ones((batch_size, kv_seq_length, kv_hidden_size), dtype=np.float32) * 0.1
+
+        Y, _, _, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            q_num_heads=q_num_heads,
+            kv_num_heads=kv_num_heads,
+        )
+
+        expect(
+            node,
+            inputs=[Q, K, V],
+            outputs=[Y],
+            name="test_attention_3d_transpose_verification",
+            opset_imports=[onnx.helper.make_opsetid("", 23)],
+        )
+
+    @staticmethod
+    def export_attention_4d_diff_heads_mask4d_padded_kv() -> None:
+        node = onnx.helper.make_node(
+            "Attention",
+            inputs=["Q", "K", "V", "attn_mask", "", "", "nonpad_kv_seqlen"],
+            outputs=["Y"],
+        )
+
+        Q = np.random.rand(2, 3, 4, 8).astype(np.float32)
+        K = np.random.rand(2, 3, 6, 8).astype(np.float32)
+        V = np.random.rand(2, 3, 6, 10).astype(np.float32)
+        attn_mask = np.random.rand(2, 3, 4, 4).astype(np.float32)
+        nonpad_kv_seqlen = np.array([3, 4], dtype=np.int64)
+
+        Y, _, _, _ = _compute_attention(
+            Q,
+            K,
+            V,
+            attn_mask=attn_mask,
+            nonpad_kv_seqlen=nonpad_kv_seqlen,
+        )
+
+        expect(
+            node,
+            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)],
         )