class_docstr: Class providing a method for fast attention via low rank decomposition.
  
  Class is responsible for providing a method <dot_product_attention> for fast
  dot-product attention with the use of low rank decomposition (e.g. with
  random feature maps).
class_docstr_tok: ["Class","providing","a","method","for","fast","attention","via","low","rank","decomposition",".","Class","is","responsible","for","providing","a","method","<","dot_product_attention",">","for","fast","dot-product","attention","with","the","use","of","low","rank","decomposition","(","e.g",".","with","random","feature","maps",")","."]
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,
  ):
      assert key.shape[:-1] == value.shape[:-1]
      assert (
          query.shape[0:1] == key.shape[0:1]
          and query.shape[-1] == key.shape[-1]
      )
      if axis is None:
          axis = tuple(range(1, key.ndim - 2))
      if not isinstance(axis, Iterable):
          axis = (axis,)
      assert key.ndim == query.ndim
      assert key.ndim == value.ndim
      for ax in axis:
          if not (query.ndim >= 3 and 1 <= ax < query.ndim - 2):
              raise ValueError(
                  "Attention axis must be between the batch "
                  "axis and the last-two axes."
              )
      n = key.ndim
  
      # Constructing projection tensor.
      if self.redraw_features:
          query_seed = lax.convert_element_type(
              jnp.ceil(jnp.sum(query) * 10000000.0), jnp.int32
          )
          rng = random.PRNGKey(query_seed)
          self.projection_matrix = self.draw_weights(rng)
  
      # batch_dims is  <bs, <non-attention dims>, num_heads>
      batch_dims = tuple(np.delete(range(n), axis + (n - 1,)))
      # q & k -> (bs, <non-attention dims>, num_heads, <attention dims>, channels)
      qk_perm = batch_dims + axis + (n - 1,)
      k_extra_perm = axis + batch_dims + (n - 1,)
      key_extra = key.transpose(k_extra_perm)
      key = key.transpose(qk_perm)
      query = query.transpose(qk_perm)
      # v -> (bs, <non-attention dims>, num_heads, <attention dims>, channels)
      v_perm = batch_dims + axis + (n - 1,)
      value = value.transpose(v_perm)
      batch_dims_t = tuple(range(len(batch_dims)))
      attention_dims_t = tuple(
          range(len(batch_dims), len(batch_dims) + len(axis))
      )
  
      # Constructing tensors Q^{'} and K^{'}.
      query_prime = self.kernel_feature_creator(
          query,
          self.projection_matrix,
          attention_dims_t,
          batch_dims_t,
          precision,
          True,
      )
      query_prime = query_prime.astype(dtype)
      key_prime = self.kernel_feature_creator(
          key,
          self.projection_matrix,
          attention_dims_t,
          batch_dims_t,
          precision,
          False,
      )
      key_prime = key_prime.astype(dtype)
  
      if self.unidirectional:
          index = attention_dims_t[0]
          z_slice_shape = (
              key_prime.shape[0 : len(batch_dims_t)]
              + (key_prime.shape[-1],)
              + (value.shape[-1],)
          )
  
          numerator_fn = _numerator(
              z_slice_shape, precision, self.lax_scan_unroll
          )
          w = numerator_fn(
              jnp.moveaxis(query_prime, index, 0),
              jnp.moveaxis(key_prime, index, 0),
              jnp.moveaxis(value, index, 0),
          )
  
          # Constructing w = (Q^{'}(K^{'})^{T})_{masked}V
          w = jnp.moveaxis(w, 0, index)
  
          if not self.renormalize_attention:
              # Unidirectional, not-normalized attention.
              perm_inv = _invert_perm(qk_perm)
              result = w.transpose(perm_inv)
              return result
          else:
              # Unidirectional, normalized attention.
              thick_all_ones = jnp.zeros(key.shape[0:-1]) + jnp.ones(
                  key_extra.shape[0 : len(axis)]
              )
  
              index = attention_dims_t[0]
              t_slice_shape = key_prime.shape[0 : len(batch_dims_t)] + (
                  key_prime.shape[-1],
              )
              denominator_fn = _denominator(
                  t_slice_shape, precision, self.lax_scan_unroll
              )
              r = denominator_fn(
                  jnp.moveaxis(query_prime, index, 0),
                  jnp.moveaxis(key_prime, index, 0),
              )
  
              r = jnp.moveaxis(r, 0, index)
      else:
          contract_query = tuple(
              range(
                  len(batch_dims) + len(axis),
                  len(batch_dims) + len(axis) + 1,
              )
          )
          contract_z = tuple(
              range(len(batch_dims), len(batch_dims) + 1)
          )
          # Constructing  z = (K^{'})^{T}V
          #  z (bs, <non-attention dims>, num_heads, channels_m, channels_v)
          z = lax.dot_general(
              key_prime,
              value,
              (
                  (attention_dims_t, attention_dims_t),
                  (batch_dims_t, batch_dims_t),
              ),
              precision=precision,
          )
          # Constructing w = Q^{'} z = Q^{'}(K^{'})^{T}V
          # q (bs, <non-attention dims>, num_heads, <attention dims>, channels_m)
          #  z (bs, <non-attention dims>, num_heads, channels_m, channels_v)
          # w (bs,  <non-attention dims>, num_heads, <attention dims>, channels_v)
          w = lax.dot_general(
              query_prime,
              z,
              (
                  (contract_query, contract_z),
                  (batch_dims_t, batch_dims_t),
              ),
              precision=precision,
          )
          if not self.renormalize_attention:
              # Bidirectional, not-normalized attention.
              perm_inv = _invert_perm(qk_perm)
              result = w.transpose(perm_inv)
              return result
          else:
              # Bidirectional, normalized attention.
              thick_all_ones = jnp.zeros(key.shape[0:-1]) + jnp.ones(
                  key_extra.shape[0 : len(axis)]
              )
              thick_all_ones = thick_all_ones.astype(dtype)
              contract_key = tuple(
                  range(len(batch_dims), len(batch_dims) + len(axis))
              )
              contract_thick_all_ones = tuple(
                  range(
                      thick_all_ones.ndim - len(axis),
                      thick_all_ones.ndim,
                  )
              )
              # Construct t = (K^{'})^{T} 1_L
              # k (bs, <non-attention dims>, num_heads, <attention dims>, channels)
              t = lax.dot_general(
                  key_prime,
                  thick_all_ones,
                  (
                      (contract_key, contract_thick_all_ones),
                      (batch_dims_t, batch_dims_t),
                  ),
                  precision=precision,
              )
  
              # Construct partition function: r = Q^{'} t = Q^{'}(K^{'})^{T} 1_L
              # q_p (bs, <non-attention dims>, num_heads, <attention dims>, channs_m)
              # t   (bs, <non-attention dims>, num_heads, channels_m)
              r = lax.dot_general(
                  query_prime,
                  t,
                  (
                      ((query_prime.ndim - 1,), (t.ndim - 1,)),
                      (batch_dims_t, range(0, len(t.shape) - 1)),
                  ),
                  precision=precision,
              )
  
      r = r + 2 * self.numerical_stabilizer * (
          jnp.abs(r) <= self.numerical_stabilizer
      )
      r = jnp.reciprocal(r)
      r = jnp.expand_dims(r, len(r.shape))
      # w (bs, <non-attention dims>, num_heads, <attention dims>, channels_v)
      # r (bs, <non-attention dims>, num_heads, <attention dims>, extra_channel)
      result = w * r
      # back to (bs, dim1, dim2, ..., dimN, num_heads, channels)
      perm_inv = _invert_perm(qk_perm)
      result = result.transpose(perm_inv)
      return result
  
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 3 for calling the FastAttentionviaLowRankDecomposition 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) and returns: result, result, result