Skip to content

Validate dropout_rate in dot_product_attention#5503

Open
mohsinm-dev wants to merge 2 commits into
google:mainfrom
mohsinm-dev:validate-dropout-rate
Open

Validate dropout_rate in dot_product_attention#5503
mohsinm-dev wants to merge 2 commits into
google:mainfrom
mohsinm-dev:validate-dropout-rate

Conversation

@mohsinm-dev

Copy link
Copy Markdown
Contributor

Fixes #5497

Reject dropout_rate outside [0, 1) (including NaN/inf) with a clear ValueError, and raise early when dropout_rng is missing for active dropout.

np.testing.assert_allclose(nnx_out, jax_out, atol=1e-3, rtol=1e-3)


class TestDropoutRateValidation(parameterized.TestCase):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reduce the number of tests to essential only and use parameterize as much as you can to reduce the number of lines

@mohsinm-dev mohsinm-dev force-pushed the validate-dropout-rate branch from 66934c0 to 89b5e6b Compare June 18, 2026 01:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

flax.nnx.dot_product_attention does not validate dropout_rate (negative/NaN silently disables dropout)

2 participants