def prng(s):
  """Creates RNG state from seed.

  Args:
    s: the seed, an integer.

  Returns:
    An RNG state, as a scalar array of dtype `np.int64`.
  """
  # TODO(wangpeng): Become bitwise-identical to JAX when TF stateless RNGs get
  #   improved.
  return tf_np.asarray(s, dtype=_RNG_KEY_DTYPE)
Creates RNG state from seed.

  Args:
    s: the seed, an integer.

  Returns:
    An RNG state, as a scalar array of dtype `np.int64`.
  
trax/tf_numpy/extensions/extensions.py