class_docstr: Relative attention layer.
  
  Layer that maps (location_bias, context_bias, pos_emb, q, k, v, mask)
      to (activations, mask).
  This layer type performs the inner workings of one pass of multi-head
  self-attention. It:
    - splits queries, keys, and values into multiple 'heads',
    - splits positional embeddings into multiple 'heads',
    - computes per-head attention weights from per-head (queries, keys),
    - applies mask to screen out positions that come from padding tokens,
    - [in `'train'` mode] applies dropout to attention weights,
    - uses attention weights to combine per-head values vectors, and
    - merges per-head results into outgoing activations matching original input
      activation vector shapes.
docstr: Returns attention-computed activations and unmodified mask.
  
  Args:
    inputs: A (location_bias, context_bias, pos_emb, q, k, v, mask) tuple.
docstr_tok: ["Returns","attention-computed","activations","and","unmodified","mask",".","Args",":","inputs",":","A","(","location_bias",",","context_bias",",","pos_emb",",","q",",","k",",","v",",","mask",")","tuple","."]
code: def forward(self, inputs):
      """Returns attention-computed activations and unmodified mask.
  
      Args:
        inputs: A (location_bias, context_bias, pos_emb, q, k, v, mask) tuple.
      """
      location_bias, context_bias, pos_emb, q, k, v, mask = inputs
  
      d_feature = q.shape[-1]
      n_heads = self._n_heads
      if d_feature % n_heads != 0:
          raise ValueError(
              f"Dimensionality of feature embedding ({d_feature}) is not a "
              f"multiple of the requested number of attention heads ({n_heads})."
          )
  
      per_head_results, dots = DotProductAttention(
          SplitIntoHeads(n_heads, merged_batch_and_head=False).forward(
              q
          ),
          SplitIntoHeads(n_heads, merged_batch_and_head=False).forward(
              k
          ),
          SplitIntoHeads(n_heads, merged_batch_and_head=False).forward(
              v
          ),
          pos_emb.reshape((-1, n_heads, d_feature // n_heads)),
          context_bias,
          location_bias,
          mask,
          separate_cls=self._separate_cls,
          dropout=self._dropout,
          mode=self._mode,
          rng=self.rng,
      )
      if self._mode == "viz":
          self.state = dots
      merged_results = MergeHeads(
          n_heads, merged_batch_and_head=False
      ).forward(per_head_results)
      return merged_results, mask
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.attention import MergeHeads
  from trax.layers.attention import SplitIntoHeads
label_desc: Use image node_id 2 for calling the RelativeAttention obj's underlying member method code with example usage: obj.forward(inputs) and returns: merged_results,  mask