code: def sinusoidal_init(max_len=2048, min_scale=1.0, max_scale=10000.0):
      """1D Sinusoidal Position Embedding Initializer.
  
      Args:
          max_len: maximum possible length for the input.
          min_scale: float: minimum frequency-scale in sine grating.
          max_scale: float: maximum frequency-scale in sine grating.
  
      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]
          scale_factor = -np.log(max_scale / min_scale) / (
              d_feature // 2 - 1
          )
          div_term = min_scale * np.exp(
              np.arange(0, d_feature // 2) * scale_factor
          )
          pe[:, : d_feature // 2] = np.sin(position * div_term)
          pe[:, d_feature // 2 : 2 * (d_feature // 2)] = np.cos(
              position * div_term
          )
          pe = pe[np.newaxis, :, :]  # [1, max_len, d_feature]
          return jnp.array(pe)
  
      return init
  
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Optional
  import jax
  import jax.numpy
  import numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import flaglib
  from flax.experimental.nnx.examples.lm1b.configs import default
label_desc: Use image node_id 3 for calling a global function with example usage: sinusoidal_init(max_len, min_scale, max_scale) and returns: init, jnp