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

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

.. flax_module::
  :module: flax.nnx
  :class: MultiHeadAttention

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