Attention
------------------------

.. automodule:: flax.nnx
.. currentmodule:: flax.nnx

.. autoclass:: MultiHeadAttention
   :members:

.. autofunction:: combine_masks
.. autofunction:: dot_product_attention
.. autofunction:: make_attention_mask
.. autofunction:: make_causal_mask