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