def _get_noise(self, lo: int, hi: int, depth: int):
    """Return pseudorandom noise with shape float[length, depth].

    Args:
      lo: where to start sampling
      hi: where to stop sampling
      depth: noise depth

    Returns:
      noise[lo:hi, :]: the noise, where noise.diff(axis=0) is i.i.d. U(-1,1)
    """
    if self._noise is None or self._noise.shape[0] < hi:
      # Resize the noise:
      new_length = 1
      while new_length < hi:
        new_length *= 2
      noise = threefry_2x32_prange(self._noise_rng, 0, new_length * depth)
      noise = noise.reshape((new_length, depth, 2))[:, :, 0]
      # Normalize to [-sqrt(3), sqrt(3)]:
      noise = noise.astype(jnp.float32) / np.float32(2**31 - 1)
      noise = noise * 3**.5
      # TODO(tying): use multiscale noise for memory-efficient sampling
      noise = noise.cumsum(axis=0)
      self._noise = noise
    assert self._noise.shape[0] >= hi
    assert self._noise.shape[1] == depth
    return self._noise[lo:hi, :]
Return pseudorandom noise with shape float[length, depth].

    Args:
      lo: where to start sampling
      hi: where to stop sampling
      depth: noise depth

    Returns:
      noise[lo:hi, :]: the noise, where noise.diff(axis=0) is i.i.d. U(-1,1)
    
trax/layers/research/position_encodings.py