Validate dropout_rate in dot_product_attention

flax/nnx/nn/attention.py

@@ -115,6 +115,15 @@ def dot_product_attention_weights(
   Returns:
     Output of shape `[batch..., num_heads, q_length, kv_length]`.
   """
+  if not 0.0 <= dropout_rate < 1.0:
+    raise ValueError(
+      f'dropout_rate must be in the range [0, 1), got {dropout_rate!r}.'
+    )
+  if not deterministic and dropout_rate > 0.0 and dropout_rng is None:
+    raise ValueError(
+      'dropout_rng is required when dropout_rate > 0 and deterministic=False.'
+    )
+
   query, key = promote_dtype((query, key), dtype=dtype)  # type: ignore[bad-unpacking]
   dtype = query.dtype
 

tests/nnx/nn/attention_test.py

@@ -383,6 +383,23 @@ def _run(m):
     np.testing.assert_allclose(nnx_out, jax_out, atol=1e-3, rtol=1e-3)
 
 
+class TestDropoutRateValidation(parameterized.TestCase):
+
+  @parameterized.parameters(
+    (-0.5, jax.random.key(1), r'dropout_rate must be in the range \[0, 1\)'),
+    (1.0, jax.random.key(1), r'dropout_rate must be in the range \[0, 1\)'),
+    (1.5, jax.random.key(1), r'dropout_rate must be in the range \[0, 1\)'),
+    (float('nan'), jax.random.key(1), r'dropout_rate must be in the range \[0, 1\)'),
+    (float('inf'), jax.random.key(1), r'dropout_rate must be in the range \[0, 1\)'),
+    (0.5, None, r'dropout_rng is required'),
+  )
+  def test_invalid_dropout_args(self, rate, rng, match):
+    q = jnp.ones((2, 4, 2, 8))
+    with self.assertRaisesRegex(ValueError, match):
+      nnx.dot_product_attention(q, q, q, dropout_rate=rate,
+                                deterministic=False, dropout_rng=rng)
+
+
 class TestDotProductAttentionValidation(parameterized.TestCase):
 
   @parameterized.parameters(