better error message

xadupre · xadupre · commit ce6ba0b24af6 · 2025-11-09T15:13:20.000+01:00
diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py
@@ -1840,7 +1840,8 @@ def aten_scaled_dot_product_attention(
         key, value = _attention_repeat_kv_for_group_query(query, key, value)
     else:
         assert query.shape[1] == key.shape[1] == value.shape[1], (
-            "SDPA (MHA) requires q_num_heads = kv_num_heads"
+            "SDPA (MHA) requires q_num_heads = kv_num_heads, "
+            f"query.shape={query.shape}, key.shape{key.shape}, value.shape={value.shape}"
         )
 
     if attn_mask is None:

Original file line number	Diff line number	Diff line change
`@@ -1840,7 +1840,8 @@ def aten_scaled_dot_product_attention(`
`1840`	`1840`	`key, value = _attention_repeat_kv_for_group_query(query, key, value)`
`1841`	`1841`	`else:`
`1842`	`1842`	`assert query.shape[1] == key.shape[1] == value.shape[1], (`
`1843`		`- "SDPA (MHA) requires q_num_heads = kv_num_heads"`
	`1843`	`+ "SDPA (MHA) requires q_num_heads = kv_num_heads, "`
	`1844`	`+ f"query.shape={query.shape}, key.shape{key.shape}, value.shape={value.shape}"`
`1844`	`1845`	`)`
`1845`	`1846`
`1846`	`1847`	`if attn_mask is None:`