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**.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(.5, .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
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]
    
trax/layers/research/position_encodings.py