docs(attention): clarify make_causal_mask and make_attention_mask docstrings (#1520)

[Ultron09]

What does this PR do?

Fixes #1520

The make_causal_mask and make_attention_mask functions generate masks of shape [batch..., 1, len_q, len_kv] for multi-head attention. However, users frequently pass embedded feature arrays ([batch..., len, num_features]) instead of sequence arrays ([batch..., len]), causing the sequence length to erroneously bind to num_features.

This PR:

1. Expands docstrings in both flax.linen.attention and flax.nnx.nn.attention with explicit shape explanations, notes on un-embedded slices vs embedded feature inputs, and how the singleton dimension broadcasts across attention heads.
2. Adds executable doctests and examples to make_attention_mask and make_causal_mask for both token IDs and embedded inputs.

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.
• Tested with doctest (passed all 90 attempted tests across Linen & NNX attention modules).

cc @marcvanzee @vfdev-5 for review.