class_docstr: Multi-head dot-product self-attention with reltive positional encodeing.
  
  Attributes:
    num_heads: Number of attention heads.
class_docstr_tok: ["Multi-head","dot-product","self-attention","with","reltive","positional","encodeing",".","Attributes",":","num_heads",":","Number","of","attention","heads","."]
docstr: Applies multi-head dot product self-attention on the input data.
  
  Args:
    x: input of shape `[bs, len, features]`.
  
  Returns:
    output of shape `[bs, len, features]`.
docstr_tok: ["Applies","multi-head","dot","product","self-attention","on","the","input","data",".","Args",":","x",":","input","of","shape","`","[","bs",",","len",",","features","]","`",".","Returns",":","output","of","shape","`","[","bs",",","len",",","features","]","`","."]
code: def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
      """Applies multi-head dot product self-attention on the input data.
  
      Args:
        x: input of shape `[bs, len, features]`.
  
      Returns:
        output of shape `[bs, len, features]`.
      """
      bs, input_len, features = x.shape
      if features % (2 * self.num_heads) != 0:
          raise ValueError(
              f"Inputs feature dimension {features} must be divisible by 2 * number "
              f"of heads {2*self.num_heads}."
          )
      head_features = features // self.num_heads
  
      # Apply qkv transformation.
      qkv = nn.Dense(
          features=features * 2,
          kernel_init=nn.initializers.normal(stddev=1.0 / features),
          use_bias=False,
          dtype=jnp.float32,
      )(x)
  
      # Normalize.
      qkv = nn.GroupNorm(epsilon=1e-4, name="gn_qkv")(qkv)
      # Split to multi-headed.
      qkv = qkv.reshape(
          bs, input_len, self.num_heads, head_features * 2
      ).transpose(
          (0, 2, 3, 1)
      )  # To bhdl.
      # Following the reference implementation, we set feature-size per head
      #  of query, key, to half of the feature size per head for value:
      query, key, value = jnp.split(
          qkv, [head_features // 2, head_features], axis=2
      )
  
      # Compute relative positional attention logits.
      length = query.shape[-1]  # Shape: `[bs, heads, depth, len]`.
      relative_emb = self.param(
          "relative_pos_emb",
          nn.initializers.normal(stddev=1.0 / head_features),
          (head_features * 2, 2 * length - 1),
          jnp.float32,
      )
      relative_pos_emb = jnp.take(
          relative_emb,
          nn_ops.compute_1d_relative_distance(length, length),
          axis=-1,
      )
      (
          relative_pos_emb_q,
          relative_pos_emb_k,
          relative_pos_emb_v,
      ) = jnp.split(
          relative_pos_emb, [head_features // 2, head_features], axis=0
      )
  
      # When computing the similarity of keys and queries, we in fact have
      # (Q + P_q)*(K + P_k), which is Q*K + K*P_q + Q*P_k + P_q*P_k.
      # The last term, i.e. P_q*P_k, is fixed for all Q, K, but but we compute
      # the attention logits for other three terms:
      # 1) We attend from content of queries to the relative position of keys,
      # i.e. Q*P_k:
      qr_attn_logits = jnp.einsum(
          "bhdi,dij->bhij", query, relative_pos_emb_k
      )
      # 2) We attend from content of keys to the relative position of queries,
      # i.e. K*P_q
      kr_attn_logits = jnp.einsum(
          "bhdi,dij->bhij", key, relative_pos_emb_q
      )
      # 3) We attend from content of queries to the content of keys, i.e. Q*K:
      qk_attn_logits = jnp.einsum("bhdi,bhdj->bhij", query, key)
      # Finally we combine all these attention logits:
      attn_weights = qk_attn_logits + qr_attn_logits + kr_attn_logits
  
      # Normalize the attention weights with softmax.
      attn_weights = jax.nn.softmax(attn_weights, axis=3).astype(
          jnp.float32
      )
      # Weighted sum over values for each query position.
      wv = jnp.einsum("bhij,bhdj->bhdi", attn_weights, value)
      wve = jnp.einsum(
          "bhij,dij->bhdi", attn_weights, relative_pos_emb_v
      )
      # From bhdi to bihd and smash head dim.
  
      x = (wv + wve).transpose((0, 3, 1, 2))  # From bhdi to bihd.
      # Smash head dim and back to the original inputs dimensions.
      out = x.reshape(bs, input_len, features)
      return nn.GroupNorm(epsilon=1e-4, name="gn_out")(out)
  
imports: from typing import Dict, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.model_lib.layers import nn_ops
  from scenic.projects.baselines import bit_resnet
label_desc: Use image node_id 1 for calling the SelfAttentionWith1DRelativePos obj's underlying member method code with example usage: obj.__call__(x) and returns: nn