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:
          # TODO(kchoro): Get rid of the constant below.
          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(onp.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,
      )
      key_prime = self.kernel_feature_creator(
          key,
          self.projection_matrix,
          attention_dims_t,
          batch_dims_t,
          precision,
          False,
      )
  
      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)]
              )
              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 functools
  from collections.abc import Iterable
  import numpy
  from absl import logging
  import jax
  import jax.numpy
  from jax import lax, random
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