def CausalAttention(d_feature,
                    n_heads=1,
                    dropout=0.0,
                    max_inference_length=2048,
                    use_dconv=False,
                    mode='train'):
  """Returns a layer that maps activations to activations, with causal masking.

  Like :py:class:`Attention`, this layer type represents one pass of multi-head
  self-attention, but with causal masking rather than padding-based masking.

  Args:
    d_feature: Last/innermost dimension of activations in the input to and
        output from this layer.
    n_heads: Number of attention heads. Attention heads effectively split
        activation vectors into ``n_heads`` subvectors, of size
        ``d_feature / n_heads``.
    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.
    use_dconv: if True, use depthwise convolutions on top of dense layers
      for Q, K and V.
    mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
  """
  if d_feature % n_heads != 0:
    raise ValueError(
        f'Dimensionality of feature embedding ({d_feature}) is not a multiple '
        f'of the requested number of attention heads ({n_heads}).')

  def QKVLayer():
    """Function returning the Q, K and V layer."""
    if use_dconv:
      return cb.Serial(core.Dense(d_feature), convolution.CausalDepthwiseConv())
    else:
      return core.Dense(d_feature)

  return ConfigurableAttention(
      QKVLayer(),
      QKVLayer(),
      QKVLayer(),
      core.Dense(d_feature),
      n_heads=n_heads,
      qkv_attention_layer=DotProductCausalAttention(
          dropout=dropout, max_inference_length=max_inference_length,
          mode=mode))
Returns a layer that maps activations to activations, with causal masking.

  Like :py:class:`Attention`, this layer type represents one pass of multi-head
  self-attention, but with causal masking rather than padding-based masking.

  Args:
    d_feature: Last/innermost dimension of activations in the input to and
        output from this layer.
    n_heads: Number of attention heads. Attention heads effectively split
        activation vectors into ``n_heads`` subvectors, of size
        ``d_feature / n_heads``.
    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.
    use_dconv: if True, use depthwise convolutions on top of dense layers
      for Q, K and V.
    mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
  
trax/layers/attention.py