def threefry_2x32_prange(key, lo: int = 0, hi: int = 2):
  """Splits a key into a stream of random keys.

  This uses the little-endian counter mode.

  Args:
    key: uint32[2] the key to split
    lo: the range to start extracting from
    hi: the range to stop extracting from

  Returns:
    keys: uint32[hi - lo, 2] the split keys
  """
  if not (key.shape == (2,) and key.dtype == jnp.uint32):
    raise ValueError('key must be uint32[2]')
  if not hi < 2**32:
    # You shouldn't really be using more than half the key size anyways.
    raise NotImplementedError('only 32-bit sizes are supported')
  # Create a 64-bit counter:
  i_lo = jnp.arange(lo, hi, dtype=jnp.uint32)
  i_hi = jnp.zeros_like(i_lo)
  i = jnp.stack([i_lo, i_hi], axis=-1)
  return threefry_2x32_prf(key, i)
Splits a key into a stream of random keys.

  This uses the little-endian counter mode.

  Args:
    key: uint32[2] the key to split
    lo: the range to start extracting from
    hi: the range to stop extracting from

  Returns:
    keys: uint32[hi - lo, 2] the split keys
  
trax/layers/research/position_encodings.py