def __init__(self, dropout=0.0, max_inference_length=2048, mode='train'):
    """Creates a :py:class:`DotProductCausalAttention` instance.

    Args:
      dropout: Probababilistic rate for attention dropout, which overrides
          (sets to zero) some attention strengths derived from query-key
          matching. As a result, on a given forward pass, some value vectors
          don't contribute to the output, analogous to how regular dropout can
          cause some node activations to be ignored. Applies only if layer is
          created in ``'train'`` mode.
      max_inference_length: Maximum sequence length allowed in non-training
          modes.
      mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
    """
    super().__init__(n_in=3, n_out=1)
    self._dropout = dropout
    self._mode = mode
    self._max_len = max_inference_length
    self._portal_mask = self.monkey_patched_mask()  # pylint: disable=assignment-from-none
Creates a :py:class:`DotProductCausalAttention` instance.

    Args:
      dropout: Probababilistic rate for attention dropout, which overrides
          (sets to zero) some attention strengths derived from query-key
          matching. As a result, on a given forward pass, some value vectors
          don't contribute to the output, analogous to how regular dropout can
          cause some node activations to be ignored. Applies only if layer is
          created in ``'train'`` mode.
      max_inference_length: Maximum sequence length allowed in non-training
          modes.
      mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
    
trax/layers/attention.py