def forward(self, inputs):
    """Returns attention-computed activations and unmodified mask.

    Args:
      inputs: A `(Q, K, V, mask)` tuple, whose query, key, and value
          activations have not yet been subdivided into heads.
    """
    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 = _per_head_attention(
        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),
        mask,
        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)
Returns attention-computed activations and unmodified mask.

    Args:
      inputs: A `(Q, K, V, mask)` tuple, whose query, key, and value
          activations have not yet been subdivided into heads.
    
trax/layers/attention.py