Add condition to dropout and ref to isnan (#2482)

op.Dropout is only enabled when `dropout_p` is not 0, and added a
reference issue discussion about why op.Where and op.IsNaN are needed
when attention mask is boolean value.

Author: titaiwangms

onnxscript/function_libs/torch_lib/ops/nn.py

@@ -2037,7 +2037,8 @@ def _aten_scaled_dot_product_attention_no_mask_onnx(
         op.MatMul(query_scaled, key_transposed_scaled),
         axis=-1,
     )
-    attn_weight, _ = op.Dropout(attn_weight, dropout_p)
+    if dropout_p != 0:
+        attn_weight, _ = op.Dropout(attn_weight, dropout_p)
     return op.MatMul(attn_weight, value)
 
 
@@ -2080,8 +2081,10 @@ def _aten_scaled_dot_product_attention_bool_mask_onnx(
     # due to the presence of -inf in an entire row (padding tokens), resulting in 0/0 (NaN) in the softmax output.
     # This is because there's no safe/masked softmax imp in ONNX, so we need to handle NaN values explicitly to match
     # the behavior of PyTorch with boolean masks.
+    # Reference: https://github.com/pytorch/pytorch/issues/103749
     attn_weight = op.Where(op.IsNaN(attn_weight), zero, attn_weight)
-    attn_weight, _ = op.Dropout(attn_weight, dropout_p)
+    if dropout_p != 0:
+        attn_weight, _ = op.Dropout(attn_weight, dropout_p)
     return op.MatMul(attn_weight, value)
 
 
@@ -2116,7 +2119,8 @@ def _aten_scaled_dot_product_attention_float_mask_onnx(
         op.Add(op.MatMul(query_scaled, key_transposed_scaled), attn_mask),
         axis=-1,
     )
-    attn_weight, _ = op.Dropout(attn_weight, dropout_p)
+    if dropout_p != 0:
+        attn_weight, _ = op.Dropout(attn_weight, dropout_p)
     return op.MatMul(attn_weight, value)