class_docstr: Infinite positional encoding.
class_docstr_tok: ["Infinite","positional","encoding","."]
docstr: Get embeddings float[length, depth].
  
  Args:
    lo: where to start sampling
    hi: where to stop sampling
    depth: embedding depth
    rng: rng for random phase
  
  Returns:
    embeddings: float[length, depth]
docstr_tok: ["Get","embeddings","float","[","length",",","depth","]",".","Args",":","lo",":","where","to","start","sampling","hi",":","where","to","stop","sampling","depth",":","embedding","depth","rng",":","rng","for","random","phase","Returns",":","embeddings",":","float","[","length",",","depth","]"]
code: def _get_embeddings(self, lo: int, hi: int, depth, rng=None):
      """Get embeddings float[length, depth].
  
      Args:
        lo: where to start sampling
        hi: where to stop sampling
        depth: embedding depth
        rng: rng for random phase
  
      Returns:
        embeddings: float[length, depth]
      """
      noise = self._get_noise(lo, hi, (depth + 1) // 2)
      # Make the stddev around 1 after 1/drift.
      noise = noise * self._drift**0.5
  
      t, c = np.mgrid[lo:hi, :depth]
      # Make even channels cos, odd channels sin:
      c_div_2, c_mod_2 = divmod(c, 2)
      # Off-by-one correction for odd depth:
      drift = self._drift
      if depth > 2:
          drift = drift ** (((depth + 1) // 2) / (depth // 2))
      # Spend roughly half the frequencies on noise:
      freq = jnp.geomspace(0.5, 0.5 * drift**2, num=(depth + 1) // 2)[
          c_div_2
      ]
      cycles = c_mod_2 / 4 + freq * t + noise[:, c_div_2[0, :]] / 4
      assert cycles.shape == (hi - lo, depth), cycles.shape
  
      # Get random phases:
      if self._affine:
          assert rng is not None
          cycles = cycles + trax.fastmath.random.uniform(
              rng,
              (
                  1,
                  depth,
              ),
              minval=0,
              maxval=1,
          )
  
      # Convert from cycles to radians:
      embeddings = jnp.cos(jnp.pi * 2 * cycles)
  
      # Set the last channels to the time bin features:
      if self._time_bin_length is not None:
          inter_bin_idx, intra_bin_idx = divmod(
              t[:, -1:], self._time_bin_length
          )
          bin_parity = inter_bin_idx % 2
          bin_fraction = intra_bin_idx / self._time_bin_length
          embeddings = jnp.concatenate(
              [
                  embeddings[:, :-3],
                  1 / (1 + inter_bin_idx),
                  bin_fraction,
                  bin_parity.astype(jnp.float32),
              ],
              -1,
          )
  
      assert embeddings.shape == (hi - lo, depth), embeddings.shape
      return embeddings
  
imports: import logging
  import jax
  import jax.extend
  import numpy
  import trax
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
label_desc: Use image node_id 3 for calling the InfinitePositionalEncoding obj's underlying member method code with example usage: obj._get_embeddings(lo, hi, depth, rng) and returns: embeddings