class_docstr: Abstract class providing a method for fast attention.
  
  Class is responsible for providing a method <dot_product_attention> for fast
  approximate attention.
class_docstr_tok: ["Abstract","class","providing","a","method","for","fast","attention",".","Class","is","responsible","for","providing","a","method","<","dot_product_attention",">","for","fast","approximate","attention","."]
docstr: 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].
code: 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")
  
imports: import abc
  import enum
  import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Iterable
  from absl import logging
  from flax import linen
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import attention_layers
label_desc: Use image node_id 1 for calling the FastAttention obj's underlying member method code with example usage: obj.dot_product_attention(query, key, value, dtype, bias, axis, broadcast_dropout, dropout_rng, dropout_rate, deterministic, precision) without return types