code: 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,
          ),
      )
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.attention import MergeHeads
  from trax.layers.attention import SplitIntoHeads
label_desc: Use image node_id 6 for calling a global function with example usage: PositionalEmbeddings(d_feature, separate_cls, total_kv_pooling) and returns: cb, positions, pos_emb