fix(attention): validate dropout_rate range and require dropout_rng w…

[Vitaliy-Pikalo]

In dot_product_attention_weights, a negative or NaN dropout_rate silently passes the dropout_rate > 0.0 guard, causing dropout to be skipped without any warning or error — the user believes regularization is applied but it isn't.

Additionally, when dropout_rate > 0.0 and deterministic=False but dropout_rng=None, the error surfaces deep inside jax.random.bernoulli as TypeError: unexpected PRNG key type NoneType instead of a helpful message pointing to the missing argument.

Fix: add two early validation checks at the top of dot_product_attention_weights:

1. ValueError if dropout_rate is not in [0, 1)
2. ValueError if dropout will be applied but dropout_rng is None

Fixes #5497.

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).
• This change is discussed in a GitHub issue(#5497).
• The documentation and docstrings adhere to the documentation guidelines.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.