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)
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.
    
trax/layers/research/rel_attention.py