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

    Args:
      inputs: A (queries, keys, values) tuple.
    """
    q, k, v = inputs

    if self._portal_mask is not None:
      mask_for_predict = self._portal_mask.get_value()
    else:
      mask_for_predict = None

    if self._mode == 'predict':
      self.state, mask = _fast_inference_update_state(
          inputs, self.state,
          mask_for_predict=mask_for_predict)
      if self._portal_mask is not None:
        (_, k, v, _) = self.state
      else:
        (k, v, _) = self.state
    else:
      sequence_length = q.shape[-2]
      mask = _causal_mask(sequence_length)

    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 activations.

    Args:
      inputs: A (queries, keys, values) tuple.
    
trax/layers/attention.py