def __init__(self, min_val=0.0, max_val=1.0, shape=(), dtype=jnp.float32,
               sync=False):
    """Layer returning a tensor with random values distributed uniformly.

    Args:
      min_val: Lower end of uniform distribution.
      max_val: Upper end of uniform distribution.
      shape: Shape of the tensor to return. Values are sampled independently.
      dtype: Type of value to return.
      sync: Whether to synchronise `rng` across devices.
    """
    super().__init__(n_in=0, n_out=1)
    self._min_val = min_val
    self._max_val = max_val
    self._shape = shape
    self._dtype = dtype
    self._sync = sync
Layer returning a tensor with random values distributed uniformly.

    Args:
      min_val: Lower end of uniform distribution.
      max_val: Upper end of uniform distribution.
      shape: Shape of the tensor to return. Values are sampled independently.
      dtype: Type of value to return.
      sync: Whether to synchronise `rng` across devices.
    
trax/layers/core.py