def CreateAttentionMaskLayer():
  """Creates attention mask layer.

  Returns a layer that based on queries, keys and accumulated pool size of
  keys/values until this layer calculates positional embeddings for
  causal relative attention calculations.

  Takes as input q, k, v and appends proper mask in the end.
  Causal attention uses masking to prevent a given sequence position from
  attending to positions greater than / following it. This is used, for
  example, when training autoregressive sequence models, or when decoding a
  sequence symbol by symbol.

  Returns:
    an attention mask layer.
  """

  def calculate_mask(queries, keys):
    batch_size = queries.shape[0]
    keys_len, queries_len = keys.shape[-2], queries.shape[-2]
    funnel_factor, is_upsampling = calc_funnel_ratio(keys_len, queries_len)

    return _funnel_mask(batch_size, keys_len, queries_len, funnel_factor,
                        is_upsampling)

  def _funnel_mask(batch_size, keys_len, queries_len, funnel_factor,
                   is_upsampling):
    """Funnel mask.

    Args:
      batch_size: batch size.
      keys_len: keys length.
      queries_len: queries length.
      funnel_factor: funnel factor.
      is_upsampling: True or False.

    Returns:
      funnel mask.

    This function based on keys/queries lengths creates a triangle mask
    that prevents tokens from attending to positions following it.

    If funnel_factor is not equal to 1 due to funnel upsampling or
    downsampling it adjusts created mask for funnel attention
    by repeating each element funnel_factor times.

    This is because after funnel layer one token attends to funnel_factor
    different tokens in downsampling. During upsampling on the other hand
    funnel_factor tokens are attending to single token before upsampling.
    """

    if funnel_factor != 1:
      if not is_upsampling:
        mask = jnp.tril(jnp.ones((queries_len, queries_len), dtype=jnp.bool_))
        mask = jnp.repeat(mask, funnel_factor, axis=-1)
      else:
        mask = jnp.tril(jnp.ones((keys_len, keys_len), dtype=jnp.bool_))
        mask = jnp.repeat(mask, funnel_factor, axis=-2)
    else:
      mask = jnp.tril(jnp.ones((queries_len, queries_len), dtype=jnp.bool_))

    return jnp.repeat(mask[None, None, :, :], batch_size, axis=0)

  return cb.Branch(
      cb.Select([0]), cb.Select([1]), cb.Select([2]),
      cb.Fn('create attention mask layer', calculate_mask, n_out=1))
Creates attention mask layer.

  Returns a layer that based on queries, keys and accumulated pool size of
  keys/values until this layer calculates positional embeddings for
  causal relative attention calculations.

  Takes as input q, k, v and appends proper mask in the end.
  Causal attention uses masking to prevent a given sequence position from
  attending to positions greater than / following it. This is used, for
  example, when training autoregressive sequence models, or when decoding a
  sequence symbol by symbol.

  Returns:
    an attention mask layer.
  
trax/layers/research/rel_attention.py