(d_feature, separate_cls, total_kv_pooling)
return cb.Serial(
      cb.Fn('Generate positions vectors', PositionsVectors, n_out=1),
      cb.Fn(
          'Transform to sinusoidal encodings', Sinusoidal_Embeddings, n_out=1))
Positional embeddings.

  Args:
    d_feature: Depth/dimensionality of feature embedding.
    separate_cls: True/False if we separate_cls in calculations.
    total_kv_pooling: Accumulated pool size of keys/values until this layer.

  Returns:
    a layer that based on queries, keys and accumulated pool size of
    keys/values until this layer calculates sinusoidal positional embeddings
    for relative attention calculations.
Positional embeddings.
Positional embeddings .
def PositionalEmbeddings(d_feature, separate_cls, total_kv_pooling):
  """Positional embeddings.

  Args:
    d_feature: Depth/dimensionality of feature embedding.
    separate_cls: True/False if we separate_cls in calculations.
    total_kv_pooling: Accumulated pool size of keys/values until this layer.

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

  def PositionsVectors(queries, keys):
    assert not separate_cls

    keys_len, queries_len = keys.shape[-2], queries.shape[-2]
    funnel_factor, is_upsampling = calc_funnel_ratio(keys_len, queries_len)

    if funnel_factor == 1:
      offset = keys_len - 1
      positions = (jnp.arange(keys_len) - offset) * total_kv_pooling
    else:
      if is_upsampling:
        positions = jnp.arange(-queries_len + 1, queries_len, 1.0)
      else:
        positions = jnp.arange(-keys_len + 1, keys_len, 1.0) * total_kv_pooling

    return positions

  def Sinusoidal_Embeddings(positions):
    inv_freq = 1 / (10000**(jnp.arange(0.0, d_feature, 2.0) / d_feature))
    sinusoid_freq = jnp.einsum('i,j->ij', positions, inv_freq)
    pos_emb = jnp.concatenate(
        [jnp.sin(sinusoid_freq), jnp.cos(sinusoid_freq)], axis=1)
    return pos_emb

  return cb.Serial(
      cb.Fn('Generate positions vectors', PositionsVectors, n_out=1),
      cb.Fn(
          'Transform to sinusoidal encodings', Sinusoidal_Embeddings, n_out=1))
def PositionalEmbeddings ( d_feature , separate_cls , total_kv_pooling ) : def PositionsVectors ( queries , keys ) : assert not separate_cls keys_len , queries_len = keys . shape [ - 2 ] , queries . shape [ - 2 ] funnel_factor , is_upsampling = calc_funnel_ratio ( keys_len , queries_len ) if funnel_factor == 1 : offset = keys_len - 1 positions = ( jnp . arange ( keys_len ) - offset ) * total_kv_pooling else : if is_upsampling : positions = jnp . arange ( - queries_len + 1 , queries_len , 1.0 ) else : positions = jnp . arange ( - keys_len + 1 , keys_len , 1.0 ) * total_kv_pooling return positions def Sinusoidal_Embeddings ( positions ) : inv_freq = 1 / ( 10000 ** ( jnp . arange ( 0.0 , d_feature , 2.0 ) / d_feature ) ) sinusoid_freq = jnp . einsum ( 'i,j->ij' , positions , inv_freq ) pos_emb = jnp . concatenate ( [ jnp . sin ( sinusoid_freq ) , jnp . cos ( sinusoid_freq ) ] , axis = 1 ) return pos_emb return cb . Serial ( cb . Fn ( 'Generate positions vectors' , PositionsVectors , n_out = 1 ) , cb . Fn ( 'Transform to sinusoidal encodings' , Sinusoidal_Embeddings , n_out = 1 ) )