code: def sine_table(
      features, length, min_timescale=1.0, max_timescale=10000.0
  ):
      fraction = (
          jnp.arange(0, features, 2, dtype=jnp.float32) / features
      )
      timescale = (
          min_timescale * (max_timescale / min_timescale) ** fraction
      )
      rotational_frequency = 1.0 / timescale
      # Must use high precision einsum here, bfloat16 rounding is catastrophic.
      sinusoid_inp = jnp.einsum(
          "i,j->ij",
          jnp.arange(length),
          rotational_frequency,
          precision=jax.lax.Precision.HIGHEST,
      )
      sinusoid_inp = jnp.concatenate(
          [sinusoid_inp, sinusoid_inp], axis=-1
      )
      return jnp.sin(sinusoid_inp), jnp.cos(sinusoid_inp)
  
imports: import dataclasses
  import typing
  import jax
  import jax.numpy
  import numpy
  from jax.sharding import PartitionSpec
  from flax.experimental import nnx
label_desc: Use image node_id 4 for calling a global function with example usage: sine_table(features, length, min_timescale, max_timescale) and returns: jnp,  jnp