def get_prng(seed):
  """JAX-compatible way of getting PRNG seeds."""
  if np.shape(seed):
    raise TypeError('PRNGKey seed must be a scalar.')
  convert = lambda k: np.reshape(np.asarray(k, np.uint32), [1])
  k1 = convert(np.bitwise_and(np.right_shift(seed, 32), 0xFFFFFFFF))
  k2 = convert(np.bitwise_and(seed, 0xFFFFFFFF))
  return np.concatenate([k1, k2], 0)
JAX-compatible way of getting PRNG seeds.
trax/fastmath/numpy.py