code: def sinusoidal_init(max_len: int, max_timescale: float = 1.0e4):
      """1D Sinusoidal Position Embedding Initializer.
  
      Args:
          max_len: maximum possible length for the input.
          max_timescale: Maximum time scale.
  
      Returns:
          output: init function returning `(1, max_len, d_feature)`
      """
  
      def init(
          key: jnp.ndarray,
          shape: Sequence[int],
          dtype: jnp.dtype = jnp.float32,
      ) -> jnp.ndarray:
          """Sinusoidal init.
  
          The defined API by JAX for a custom initializer is:
            `def init(key, shape, dtype)`
  
          Even though some of args might be not used, the signature should follow
          this API as JAX passes all the three arguments (key, shape, dtype)
          to the initializers.
  
          Args:
            key: JAXPRNG key.
            shape: Shape used for making the initialized values.
            dtype: JAX data type.
  
          Returns:
            Initialized values
          """
          del key, dtype
          d_feature = shape[-1]
          pos_emb = np.zeros((max_len, d_feature), dtype=np.float32)
          position = np.arange(0, max_len)[:, np.newaxis]
          div_term = np.exp(
              np.arange(0, d_feature, 2)
              * -(np.log(max_timescale) / d_feature)
          )
          pos_emb[:, 0::2] = np.sin(position * div_term)
          pos_emb[:, 1::2] = np.cos(position * div_term)
          pe = pos_emb[
              np.newaxis, :, :
          ]  #  Shape: `[1, max_len, d_feature]`.
          return jnp.array(pe)
  
      return init
  
imports: import functools
  from typing import Callable, Optional, Sequence, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 4 for calling a global function with example usage: sinusoidal_init(max_len, max_timescale) and returns: init, jnp