examples/research_projects/performer/modeling_flax_performer_utils.py/FastAttention/dot_product_attention
class FastAttention:    def dot_product_attention(
        self,
        query,
        key,
        value,
        dtype=jnp.float32,
        bias=None,
        axis=None,
        broadcast_dropout=True,
        dropout_rng=None,
        dropout_rate=0.0,
        deterministic=False,
        precision=None,
    ):
        """
        Computes dot-product attention given query, key, and value. This is the core function for applying fast
        approximate dot-product attention. It calculates the attention weights given query and key and combines the
        values using the attention weights. This function supports multi-dimensional inputs

        Args:
          query: queries for calculating attention with shape of [batch_size, dim1,
            dim2, ..., dimN, num_heads, mem_channels].
          key: keys for calculating attention with shape of [batch_size, dim1, dim2,
            ..., dimN, num_heads, mem_channels].
          value: values to be used in attention with shape of [batch_size, dim1,
            dim2,..., dimN, num_heads, value_channels].
          dtype: the dtype of the computation (default: float32)
          bias: bias for the attention weights. This can be used for incorporating
            autoregressive mask, padding mask, proximity bias.
          axis: axises over which the attention is applied.
          broadcast_dropout: bool: use a broadcasted dropout along batch dims.
          dropout_rng: JAX PRNGKey: to be used for dropout.
          dropout_rate: dropout rate.
          deterministic: bool, deterministic or not (to apply dropout).
          precision: numerical precision of the computation see `jax.lax.Precision`
            for details

        Returns:
          Output of shape [bs, dim1, dim2, ..., dimN,, num_heads, value_channels].
        """
        raise NotImplementedError("Abstract method")