Add parse tests for multi head attention

ROCm · Dec 18, 2024 · 3fbb3fc · 3fbb3fc
1 parent 3394d29
commit 3fbb3fc
Show file tree

Hide file tree

Showing 17 changed files with 701 additions and 0 deletions.
diff --git a/test/onnx/gen_onnx.py b/test/onnx/gen_onnx.py
@@ -7864,6 +7864,159 @@ def mha_scale_test():
     return ([node], [query, key, value], [out])
 
 
+@onnx_test()
+def mha_invalid_attribute_test():
+    query = helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, 2, 4])
+    key = helper.make_tensor_value_info("k", TensorProto.FLOAT, [1, 2, 4])
+    value = helper.make_tensor_value_info("v", TensorProto.FLOAT, [1, 2, 4])
+    out = helper.make_tensor_value_info("out", TensorProto.FLOAT, [1, 2, 4])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['q', 'k', 'v'],
+                            outputs=['out'],
+                            domain='com.microsoft')
+
+    return ([node], [query, key, value], [out])
+
+
+@onnx_test()
+def mha_invalid_input_test():
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=[],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [], [])
+
+
+@onnx_test()
+def mha_invalid_query_test():
+    query = helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, 1])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['q'],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [query], [])
+
+
+@onnx_test()
+def mha_invalid_qkv_test():
+    qkv = helper.make_tensor_value_info("qkv", TensorProto.FLOAT,
+                                        [1, 1, 1, 1, 1])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['qkv'],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [qkv], [])
+
+
+@onnx_test()
+def mha_invalid_key_missing_test():
+    query = helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, 1, 1])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['q'],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [query], [])
+
+
+@onnx_test()
+def mha_invalid_key_ndim_test():
+    query = helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, 1, 1])
+    key = helper.make_tensor_value_info("k", TensorProto.FLOAT, [1, 1])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['q', 'k'],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [query, key], [])
+
+
+@onnx_test()
+def mha_invalid_kv_test():
+    query = helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, 1, 1])
+    kv = helper.make_tensor_value_info("kv", TensorProto.FLOAT, [1, 1, 1, 1, 1])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['q', 'kv'],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [query, kv], [])
+
+
+@onnx_test()
+def mha_invalid_key_test():
+    query = helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, 1, 1])
+    key = helper.make_tensor_value_info("k", TensorProto.FLOAT, [1, 1, 2])
+    value = helper.make_tensor_value_info("v", TensorProto.FLOAT, [1, 1, 1])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['q', 'k', 'v'],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [query, key, value], [])
+
+
+@onnx_test()
+def mha_invalid_value_missing_test():
+    query = helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, 1, 1])
+    key = helper.make_tensor_value_info("k", TensorProto.FLOAT, [1, 1, 1])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['q', 'k'],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [query, key], [])
+
+
+@onnx_test()
+def mha_invalid_value_test():
+    query = helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, 1, 1])
+    key = helper.make_tensor_value_info("k", TensorProto.FLOAT, [1, 1, 1])
+    value = helper.make_tensor_value_info("v", TensorProto.FLOAT, [2, 1, 1])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['q', 'k', 'v'],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [query, key, value], [])
+
+
+@onnx_test()
+def mha_invalid_value_ndim_test():
+    query = helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, 1, 1])
+    key = helper.make_tensor_value_info("k", TensorProto.FLOAT, [1, 1, 1])
+    value = helper.make_tensor_value_info("v", TensorProto.FLOAT, [1, 1, 1, 1])
+
+    node = helper.make_node('MultiHeadAttention',
+                            inputs=['q', 'k', 'v'],
+                            outputs=[],
+                            num_heads=1,
+                            domain='com.microsoft')
+
+    return ([node], [query, key, value], [])
+
+
 @onnx_test()
 def multinomial_test():
     sample_size = 13

diff --git a/test/onnx/mha_invalid_attribute_test.onnx b/test/onnx/mha_invalid_attribute_test.onnx
@@ -0,0 +1,25 @@
+	mha_invalid_attribute_test:�
+1
+q
+k
+vout"MultiHeadAttention:com.microsoftmha_invalid_attribute_testZ
+q
+
+
+
+Z
+k
+
+
+
+Z
+v
+
+
+
+b
+out
+
+
+
+B

diff --git a/test/onnx/mha_invalid_input_test.onnx b/test/onnx/mha_invalid_input_test.onnx
@@ -0,0 +1,3 @@
+	mha_invalid_input_test:O
+5"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_input_testB

diff --git a/test/onnx/mha_invalid_key_missing_test.onnx b/test/onnx/mha_invalid_key_missing_test.onnx
@@ -0,0 +1,9 @@
+	mha_invalid_key_missing_test:q
+8
+q"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_key_missing_testZ
+q
+
+
+
+B

diff --git a/test/onnx/mha_invalid_key_ndim_test.onnx b/test/onnx/mha_invalid_key_ndim_test.onnx
@@ -0,0 +1,14 @@
+	mha_invalid_key_ndim_test:�
+;
+q
+k"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_key_ndim_testZ
+q
+
+
+
+Z
+k
+
+
+B

diff --git a/test/onnx/mha_invalid_key_test.onnx b/test/onnx/mha_invalid_key_test.onnx
@@ -0,0 +1,21 @@
+	mha_invalid_key_test:�
+>
+q
+k
+v"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_key_testZ
+q
+
+
+
+Z
+k
+
+
+
+Z
+v
+
+
+
+B

diff --git a/test/onnx/mha_invalid_kv_test.onnx b/test/onnx/mha_invalid_kv_test.onnx
@@ -0,0 +1,17 @@
+	mha_invalid_kv_test:�
+<
+q
+kv"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_kv_testZ
+q
+
+
+
+Z 
+kv
+
+
+
+
+
+B

diff --git a/test/onnx/mha_invalid_qkv_test.onnx b/test/onnx/mha_invalid_qkv_test.onnx
@@ -0,0 +1,11 @@
+	mha_invalid_qkv_test:u
+:
+qkv"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_qkv_testZ!
+qkv
+
+
+
+
+
+B

diff --git a/test/onnx/mha_invalid_query_test.onnx b/test/onnx/mha_invalid_query_test.onnx
@@ -0,0 +1,8 @@
+	mha_invalid_query_test:g
+8
+q"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_query_testZ
+q
+
+
+B

diff --git a/test/onnx/mha_invalid_value_missing_test.onnx b/test/onnx/mha_invalid_value_missing_test.onnx
@@ -0,0 +1,15 @@
+	mha_invalid_value_missing_test:�
+;
+q
+k"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_value_missing_testZ
+q
+
+
+
+Z
+k
+
+
+
+B

diff --git a/test/onnx/mha_invalid_value_ndim_test.onnx b/test/onnx/mha_invalid_value_ndim_test.onnx
@@ -0,0 +1,22 @@
+	mha_invalid_value_ndim_test:�
+>
+q
+k
+v"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_value_ndim_testZ
+q
+
+
+
+Z
+k
+
+
+
+Z
+v
+
+
+
+
+B

diff --git a/test/onnx/mha_invalid_value_test.onnx b/test/onnx/mha_invalid_value_test.onnx
@@ -0,0 +1,21 @@
+	mha_invalid_value_test:�
+>
+q
+k
+v"MultiHeadAttention*
+	num_heads�:com.microsoftmha_invalid_value_testZ
+q
+
+
+
+Z
+k
+
+
+
+Z
+v
+
+
+
+B