docs: clarify shape dynamics and add example for make_causal_mask

[SakshieP]

Closes #1520

What does this PR do?

This PR enhances the documentation for the make_causal_mask utility function by clarifying its tensor shape dynamics and explaining how its output shape [batch..., 1, len, len] seamlessly broadcasts across head dimensions within MultiHeadAttention.

Changes Made:

• Updated the core docstring in flax/linen/attention.py.
• Added a clear, functional description relating the function to autoregressive language modeling.
• Provided a concrete, reproducible code example showing input and output shapes using standard doctest syntax.

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!) [Note: Dismissed as this is a pure docstring documentation enhancement]

[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.