code: def add_positional_embed(
      x: jnp.ndarray,
      feat_name: str,
      positional_embedding="sinusoidal_1d",
  ):
      """Adds positional embedding."""
  
      if x.ndim != 3:  # (batch, len, emb)
          raise ValueError(
              f"Input should be 3 dimensional. Got {x.shape}"
          )
      if positional_embedding != "sinusoidal_1d":
          raise ValueError("Only sinusoidal_1d embedding is supported!")
  
      return avmae_model_utils.add_positional_embeddings(
          x,
          positional_embedding,
          input_shape=x.shape,
          layer_name=f"posembed_{feat_name}",
      )
  
imports: import functools
  import re
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.av_mae import model_utils
  from scenic.projects.baselines import vit
  from scenic.projects.mbt import model
  from scenic.projects.vivit import model_utils
  import scipy
label_desc: Use image node_id 1 for calling a global function with example usage: add_positional_embed(x, feat_name, positional_embedding) and returns: avmae_model_utils