Skip to content

docs(attention): document None default behavior in MultiHeadAttention layers - #5561

Open
Ultron09 wants to merge 1 commit into
google:mainfrom
Ultron09:doc/mha-optional-none-defaults
Open

docs(attention): document None default behavior in MultiHeadAttention layers#5561
Ultron09 wants to merge 1 commit into
google:mainfrom
Ultron09:doc/mha-optional-none-defaults

Conversation

@Ultron09

Copy link
Copy Markdown

What does this PR do?

Fixes #4182

In MultiHeadAttention and MultiHeadDotProductAttention (in flax.linen.attention) as well as MultiHeadAttention (in flax.nnx.nn.attention), several optional arguments default to None, but their docstrings did not document the exact fallback evaluation behavior when None is provided or left as default.

This PR clarifies the None default behavior for:

  • qkv_features: defaults to inputs_q.shape[-1] (in Linen) or in_features (in NNX).
  • out_features: defaults to inputs_q.shape[-1] (in Linen) or in_features (in NNX).
  • deterministic: clarifies behavior when None vs dropout_rate > 0.0.
  • precision: clarifies that None uses the default JAX numerical precision.

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/discussion.
  • The documentation and docstrings adhere to the documentation guidelines.
  • This change includes necessary high-coverage tests. (No quality testing = no merge!)

cc @vfdev-5 @cgarciae for review.

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.

MultiHeadAttention documentation missing descriptions of None values for optional arguments

2 participants