code: def fixed_pos_embedding(x, seq_dim=0):
      dim = x.shape[-1]
      inv_freq = 1.0 / (10000 ** (np.arange(0, dim, 2) / dim))
  
      sinusoid_inp = np.einsum(
          "i , j -> i j", np.arange(x.shape[seq_dim]), inv_freq
      )
  
      return np.sin(sinusoid_inp), np.cos(sinusoid_inp)
  
imports: import haiku
  import jax
  import jax.numpy
  import numpy
  from einops import rearrange, repeat
  from mesh_transformer.util import f_psum, g_psum, maybe_shard, head_print
  from jax.experimental import PartitionSpec
  from jax.experimental.maps import thread_resources
label_desc: Use image node_id 2 for calling a global function with example usage: fixed_pos_embedding(x, seq_dim) and returns: np,  np