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

    Args:
      inputs: A `(Q, K, V, mask)` tuple, whose query, key, and value
          activations have been subdivided into heads.
    """
    q, k, v, mask = inputs
    activations, attn_strengths = _per_head_attention(
        q, k, v, mask, dropout=self._dropout, mode=self._mode, rng=self.rng)
    if self._mode == 'viz':
      self.state = attn_strengths
    return activations
Returns attention-computed per-head activations and unchanged mask.

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