fix plugs

Author: xadupre

_unittests/ut_tasks/try_export.py

@@ -212,18 +212,18 @@ def _config_reduction(config, task):
                 print(f"-- MODEL CONVERTED IN {time.perf_counter() - begin}")
                 model = onnx.load(filename, load_external_data=False)
                 if attention == "PACKED":
-                    self.assertIn("PackedMultiHeadAttention", str(model))
+                    self.assertIn('"PackedMultiHeadAttention"', str(model))
                 elif attention == "BIGMASK":
-                    self.assertNotIn("PackedMultiHeadAttention", str(model))
+                    self.assertNotIn('"PackedMultiHeadAttention"', str(model))
                     self.assertNotIn("MultiHeadAttention", str(model))
                     self.assertNotIn("Loop", {n.op_type for n in model.graph.node})
                 elif attention == "LOOPMHA":
-                    self.assertNotIn("PackedMultiHeadAttention", str(model))
-                    self.assertIn("MultiHeadAttention", str(model))
+                    self.assertNotIn('"PackedMultiHeadAttention"', str(model))
+                    self.assertIn('"MultiHeadAttention"', str(model))
                     self.assertIn("Loop", {n.op_type for n in model.graph.node})
                 elif attention == "LOOPA24":
-                    self.assertNotIn("PackedMultiHeadAttention", str(model))
-                    self.assertNotIn("MultiHeadAttention", str(model))
+                    self.assertNotIn('"PackedMultiHeadAttention"', str(model))
+                    self.assertNotIn('"MultiHeadAttention"', str(model))
                     self.assertIn("Loop", {n.op_type for n in model.graph.node})
                 else:
                     raise AssertionError(f"attention={attention!r} not expected")
@@ -257,7 +257,7 @@ def _config_reduction(config, task):
                         else ["CPUExecutionProvider"]
                     ),
                     use_ort=True,
-                    atol=0.02,
+                    atol=0.05,
                     rtol=10,
                     # ep=pt2_file,
                     expected=expected,

_unittests/ut_torch_onnx/test_sbs.py

@@ -693,21 +693,7 @@ def forward(self, query, key, value, seq_lens):
                 ks = key * mask
                 vs = value * mask
                 attn_output = qwen_sdpa_attention_loopmha_versatile(
-                    qs,
-                    ks,
-                    vs,
-                    seq_lens,
-                    0.11,
-                    16,
-                    (
-                        onnx.TensorProto.FLOAT
-                        if query.dtype == torch.float32
-                        else (
-                            onnx.TensorProto.FLOAT16
-                            if query.dtype == torch.float16
-                            else onnx.TensorProto.BFLOAT16
-                        )
-                    ),
+                    qs, ks, vs, seq_lens, 0.11, 16
                 )
                 red = attn_output.mean(dim=-1, keepdim=True)
                 return attn_output - red

onnx_diagnostic/export/onnx_plug.py

@@ -232,7 +232,7 @@ def get_function_proto(self, *args) -> onnx.FunctionProto:
         ):
             return self._function_proto_versioned[args[0]]
         try:
-            key = self.version_selector(*args)
+            key = self.version_selector(*args)  # type: ignore[misc]
         except (ValueError, AttributeError) as e:
             raise AssertionError(
                 f"Unable to select a version, fails to get a key, available="