code: def add_positional_embeddings(
      inputs: jnp.ndarray,
      posemb_type: str,
      input_shape: Optional[Iterable[int]] = None,
      layer_name: str = "posembed_input",
  ) -> jnp.ndarray:
      """Adds positional embeddings to an input sequence.
  
      Args:
        inputs: Tokens of shape [batch, num_tokens, hidden_size].
        posemb_type: The type of positional encoding. Must be one of
          {sinusoidal_1d, sinusoidal_2d, sinusoidal_3d, learned_1d}.
        input_shape: Used for "sinusoidal_2d" and "sinusoidal_3d". In this case,
          the input is reshaped to this size ie [batch, height, width, hidden_size],
          before applying the positional encodings and then reshaping back.
        layer_name: The layer name for learned embedddings.
  
      Returns:
        The input tokens with the positional encodings added. The shape is
          [batch, num_tokens, hidden_size].
      """
  
      if posemb_type == "learned_1d":
          x_posemb = vit.AddPositionEmbs(
              posemb_init=nn.initializers.normal(
                  stddev=0.02
              ),  # from BERT.
              name=layer_name,
          )(inputs)
      elif posemb_type == "learned_space_time":
          x_reshape = inputs.reshape(input_shape)
          x_posemb = AddFactorisedSpaceTimePositionEmbs(
              posemb_init_space=nn.initializers.normal(
                  stddev=0.02
              ),  # from BERT.
              posemb_init_time=nn.initializers.normal(stddev=0.02),
              name=layer_name,
          )(x_reshape)
          x_posemb = jnp.reshape(x_posemb, inputs.shape)
      elif posemb_type == "sinusoidal_1d":
          x_posemb = attention_layers.Add1DPositionEmbedding(
              posemb_init=None
          )(inputs)
      elif posemb_type in {"sinusoidal_2d", "sinusoidal_3d"}:
          x_reshape = inputs.reshape(input_shape)
          x_posemb = attention_layers.AddFixedSinCosPositionEmbedding()(
              x_reshape
          )
          x_posemb = jnp.reshape(x_posemb, inputs.shape)
      elif posemb_type == "none":
          x_posemb = inputs
      else:
          raise ValueError(
              f"Unknown positional embedding {posemb_type}"
          )
  
      return x_posemb
  
file_path: scenic/scenic/projects/av_mae/model_utils.py
imports: from typing import Callable, Iterable, Optional, Sequence
  import flax.linen
  import jax
  import jax.numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 4 for calling a global function with example usage: add_positional_embeddings(inputs, posemb_type, input_shape, layer_name) and returns: x_posemb