code: def sinusoidal_init(max_len=2048):
      """1D Sinusoidal Position Embedding Initializer.
  
      Args:
          max_len: maximum possible length for the input
  
      Returns:
          output: init function returning `(1, max_len, d_feature)`
      """
  
      def init(key, shape, dtype=np.float32):
          """Sinusoidal init."""
          del key, dtype
          d_feature = shape[-1]
          pe = 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(10000.0) / d_feature)
          )
          pe[:, 0::2] = np.sin(position * div_term)
          pe[:, 1::2] = np.cos(position * div_term)
          pe = pe[np.newaxis, :, :]  # [1, max_len, d_feature]
          return jnp.array(pe)
  
      return init
  
file_path: flax/examples/nlp_seq/models.py
imports: from typing import Callable, Any, Optional
  from flax import linen
  from flax import struct
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: sinusoidal_init(max_len) and returns: init, jnp