support mha def in (q,v) format

Author: calad0i

hls4ml/converters/keras_v3/hgq2/multi_head_attention.py

@@ -26,9 +26,13 @@ def handle(
         from hgq.layers import QEinsum
         from keras import KerasTensor
 
-        assert len(in_tensors) in (3, 4), 'MultiHead layer must have 3 (Q, K, V) or 4 (Q, K, V, M) input tensors'
+        # fmt: off
+        assert len(in_tensors) in (2, 3, 4,), (
+            'MultiHead layer must have 2 (Q, V), 3 (Q, V, K) or 4 (Q, V, K, M) input tensors'
+        )
+        # fmt: on
         assert len(out_tensors) == 1, 'Attention score output is not supported yet'
-        assert len(in_tensors) == 3, 'Mask tensor is not supported yet'
+        assert len(in_tensors) <= 3, 'Mask tensor is not supported yet'
         tensor_q, *_ = in_tensors
         tensor_O, *tensor_attn = out_tensors
         unique_name: str = layer.name
@@ -50,6 +54,8 @@ def handle(
         tensor_q = bound.arguments['query']
         tensor_k = bound.arguments['key']
         tensor_v = bound.arguments['value']
+        if tensor_k is None:
+            tensor_k = tensor_v
         tensor_q_mask = bound.arguments['query_mask']
         tensor_k_mask = bound.arguments['key_mask']
         tensor_v_mask = bound.arguments['value_mask']

test/pytest/test_hgq2_mha.py

@@ -9,7 +9,7 @@
 import numpy as np
 from hgq.config import QuantizerConfigScope
 from hgq.layers import QMultiHeadAttention
-from hgq.utils import trace_mode
+from hgq.utils import trace_minmax
 
 from hls4ml.converters import convert_from_keras_model
 
@@ -30,8 +30,7 @@ def test_hgq2_mha(strategy):
     data_k = np.random.randn(10000, 12, 7).astype(np.float32) * 3
     data = [data_q, data_v, data_k]
 
-    with trace_mode(model):
-        r_keras = model.predict(data, batch_size=1000)
+    r_keras = trace_minmax(model, data, return_results=True)
 
     model_hls = convert_from_keras_model(
         model,