def RandomUniformInitializer(lim=1.0):
  """Returns an initializer for random uniform coefficients."""
  # Make sure shape does not contain int tensors by calling int() below.
  return lambda shape, rng: random.uniform(  # pylint: disable=g-long-lambda
      rng, _PureShape(shape), jnp.float32, -lim, lim)
Returns an initializer for random uniform coefficients.
trax/layers/initializers.py