faster rng for jax

faster_rng.py

def hardware_uniform(rng_key: PRNGKey,
                     shape: Shape,
                     dtype: Dtype = np.float32,
                     minval: Array = np.float32(0),
                     maxval: Array = np.float32(1)) -> Array:
  del rng_key  # non-deterministic prng.
  minval = lax.convert_element_type(minval, dtype)
  maxval = lax.convert_element_type(maxval, dtype)
  return lax.rng_uniform(minval, maxval, shape)


# For dropout-only hardware rng.
def hardware_bernoulli(rng_key: PRNGKey,
                       p: np.ndarray = np.float32(0.5),
                       shape: Shape = None) -> Array:
  del rng_key  # non-deterministic prng.
  return lax.rng_uniform(0.0, 1.0, shape) < p


def set_hardware_rng():
  jax.random.bernoulli = hardware_bernoulli
  jax.random.uniform = hardware_uniform
  jax._src.random.uniform = hardware_uniform