()
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.
Creates attention mask layer.
Creates attention mask layer .
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))
def CreateAttentionMaskLayer ( ) : 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 ) )