class_docstr: Linformer Encoder only multi-head dot-product self-attention.
  
  Attributes:
    num_heads: Number of attention heads. Features (i.e. inputs_q.shape[-1])
      should be divisible by the number of heads.
    qkv_features: Dimension of the key, query, and value.
    out_features: Dimension of the last projection.
    broadcast_dropout: Use a broadcasted dropout along batch dims.
    dropout_rate: Dropout rate.
    kernel_init: Initializer for the kernel of the Dense layers.
    bias_init: Initializer for the bias of the Dense layers.
    bias: Whether pointwise QKVO dense transforms use bias.
    dtype: The dtype of the computation.
    precision: Numerical precision of the computation see `jax.lax.Precision`
      for details.
    low_rank_features: Low rank features.
    proj_mode: Supports "linear",  "mlp", or "cnn" projections.
    downsample: Supports downsampling query too.
    proj_configs: Configurations used in the low-rank projection.
    qk_attention_fn: A function that given multi-headed key, query, and value
      computes the attention and generates the new values.
class_docstr_tok: ["Linformer","Encoder","only","multi-head","dot-product","self-attention",".","Attributes",":","num_heads",":","Number","of","attention","heads",".","Features","(","i.e",".","inputs_q.shape","[","-1","]",")","should","be","divisible","by","the","number","of","heads",".","qkv_features",":","Dimension","of","the","key",",","query",",","and","value",".","out_features",":","Dimension","of","the","last","projection",".","broadcast_dropout",":","Use","a","broadcasted","dropout","along","batch","dims",".","dropout_rate",":","Dropout","rate",".","kernel_init",":","Initializer","for","the","kernel","of","the","Dense","layers",".","bias_init",":","Initializer","for","the","bias","of","the","Dense","layers",".","bias",":","Whether","pointwise","QKVO","dense","transforms","use","bias",".","dtype",":","The","dtype","of","the","computation",".","precision",":","Numerical","precision","of","the","computation","see","`","jax.lax.Precision","`","for","details",".","low_rank_features",":","Low","rank","features",".","proj_mode",":","Supports","``","linear","''",",","``","mlp","''",",","or","``","cnn","''","projections",".","downsample",":","Supports","downsampling","query","too",".","proj_configs",":","Configurations","used","in","the","low-rank","projection",".","qk_attention_fn",":","A","function","that","given","multi-headed","key",",","query",",","and","value","computes","the","attention","and","generates","the","new","values","."]
docstr: Applies Linformer multi-head dot product self-attention.
  
  Projects the inputs into multi-headed query, key, and value vectors,
  applies dot-product attention and project the results to an output vector.
  
  Args:
    inputs_q: Input query of shape `[batch_sizes..., length, features]`.
    inputs_kv: Input key-vale, which is ignored in linformer.
    deterministic: Whether the model is run in deterministic mode (if so, do
      not apply dropout).
  
  Returns:
    Output of shape `[batch_sizes..., length features]`.
code: def __call__(
      self,  # pytype: disable=annotation-type-mismatch  # jax-ndarray
      inputs_q: jnp.ndarray,
      inputs_kv: jnp.ndarray = None,
      *,
      deterministic: bool,
  ) -> jnp.ndarray:
      """Applies Linformer multi-head dot product self-attention.
  
      Projects the inputs into multi-headed query, key, and value vectors,
      applies dot-product attention and project the results to an output vector.
  
      Args:
        inputs_q: Input query of shape `[batch_sizes..., length, features]`.
        inputs_kv: Input key-vale, which is ignored in linformer.
        deterministic: Whether the model is run in deterministic mode (if so, do
          not apply dropout).
  
      Returns:
        Output of shape `[batch_sizes..., length features]`.
      """
      if inputs_kv is not None:
          logging.warning(
              "Ignoring inputs_kv as Linformer only supports self-attention."
          )
      x = inputs_q
      features = self.out_features or x.shape[-1]
      qkv_features = self.qkv_features or x.shape[-1]
  
      assert (
          qkv_features % self.num_heads == 0
      ), "Memory dimension must be divisible by number of heads."
      head_dim = qkv_features // self.num_heads
  
      # Project inputs_q to multi-headed q/k/v.
      # Dimensions are then [bs, dims..., n_heads, n_features_per_head].
      dense = functools.partial(
          nn.DenseGeneral,
          features=(self.num_heads, head_dim),
          axis=-1,
          kernel_init=self.kernel_init,
          bias_init=self.bias_init,
          use_bias=self.use_bias,
          precision=self.precision,
      )
  
      query, key, value = (
          dense(dtype=self.dtype, name="query")(x),
          dense(dtype=self.dtype, name="key")(x),
          dense(dtype=self.dtype, name="value")(x),
      )
  
      def _linear_low_rank_projection(
          key,
          value,
          features,
          activation=None,
          transpose=True,
          query=None,
      ):
          # By default, shared projections.
          if transpose:
              # Transpose if input is already transposed.
              key = key.transpose((0, 3, 2, 1))
              value = value.transpose((0, 3, 2, 1))
              if query is not None:
                  query = query.transpose((0, 3, 2, 1))
          dense_proj = functools.partial(
              nn.DenseGeneral,
              features=features,
              axis=-1,
              kernel_init=self.kernel_init,
              bias_init=self.bias_init,
              use_bias=None,
              precision=self.precision,
              dtype=self.dtype,
          )
          shared_dense_proj = dense_proj()
          key = shared_dense_proj(key)
          value = shared_dense_proj(value)
          if query is not None:
              query = shared_dense_proj(query)
          if activation is not None:
              key = activation(key)
              value = activation(value)
              if query is not None:
                  query = activation(query)
          if transpose:
              key = key.transpose((0, 3, 2, 1))
              value = value.transpose((0, 3, 2, 1))
              if query is not None:
                  query = query.transpose((0, 3, 2, 1))
          return key, value, query
  
      def _mlp_low_rank_projection(key, value, features):
          """MLP-based low rank projection function."""
          # Handle transpose outside (before and after linear low rank projections).
          key = key.transpose((0, 3, 2, 1))
          value = value.transpose((0, 3, 2, 1))
          for f in features[:-1]:
              key, value, _ = _linear_low_rank_projection(
                  key,
                  value,
                  features=f,
                  activation=nn.relu,
                  transpose=False,
              )
          # Don't apply activation on the last layer.
          key, value, _ = _linear_low_rank_projection(
              key,
              value,
              features=features[-1],
              activation=None,
              transpose=False,
          )
          key = key.transpose((0, 3, 2, 1))
          value = value.transpose((0, 3, 2, 1))
          return key, value
  
      if self.proj_mode == "linear":
          logging.info("Using linear low-rank projectors")
          if self.downsample:
              key, value, query = _linear_low_rank_projection(
                  key,
                  value,
                  features=self.low_rank_features,
                  transpose=True,
                  query=query,
              )
          else:
              key, value, _ = _linear_low_rank_projection(
                  key,
                  value,
                  features=self.low_rank_features,
                  transpose=True,
              )
      elif self.proj_mode == "mlp":
          # Note: do not support downsampling.
          logging.info("Using MLP low-rank projectors")
          key, value = _mlp_low_rank_projection(
              key,
              value,
              features=[self.low_rank_features, self.low_rank_features],
          )
      else:
          raise NotImplementedError(
              "This low-rank projection is not supported."
          )
  
      attention_bias = None
      dropout_rng = None
      if not deterministic and self.dropout_rate > 0.0:
          dropout_rng = self.make_rng("dropout")
  
      # Apply attention.
      x = self.qk_attention_fn(
          query,
          key,
          value,
          bias=attention_bias,
          broadcast_dropout=self.broadcast_dropout,
          dropout_rng=dropout_rng,
          dropout_rate=self.dropout_rate,
          deterministic=deterministic,
          dtype=self.dtype,
          precision=self.precision,
      )
  
      # Project back to the original inputs dimensions.
      out = nn.DenseGeneral(
          features=features,
          axis=(-2, -1),
          kernel_init=self.kernel_init,
          bias_init=self.bias_init,
          use_bias=self.use_bias,
          dtype=self.dtype,
          precision=self.precision,
          name="out",
      )(x)
      return out
  
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 LinformerEncoderAttention obj's underlying member method code with example usage: obj.__call__(inputs_q, inputs_kv) and returns: out, key,  value,  query, key,  value