def ScaledInitializer(out_dim, in_dim, scale, mode, distribution):
  """Returns an initializer that adjusts its scale based on weight shapes."""
  if scale <= 0.:
    raise ValueError('scale must be positive float, {} given'.format(scale))
  if mode not in {'fan_in', 'fan_out', 'fan_avg'}:
    raise ValueError(
        'Invalid mode argument:, {}, must be either fan_in, fan_out or fan_avg'
        .format(mode))

  def Init(shape, rng, nonreceptive_dims=None):
    """Returns random values for initializing weights of the given `shape`."""
    shape = _PureShape(shape)
    fan_in, fan_out = _GetFans(shape, out_dim, in_dim, nonreceptive_dims)
    gain = scale
    if mode == 'fan_in':
      gain /= fan_in
    elif mode == 'fan_out':
      gain /= fan_out
    elif mode == 'fan_avg':
      gain /= (fan_in + fan_out) / 2
    if distribution == 'truncated_normal':
      # constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
      stddev = jnp.sqrt(gain) / .87962566103423978
      new_weights = random.truncated_normal(rng, -2, 2, shape) * stddev
      return new_weights.astype('float32')
    elif distribution == 'normal':
      new_weights = random.normal(rng, shape) * jnp.sqrt(gain)
      return new_weights.astype('float32')
    elif distribution == 'uniform':
      lim = jnp.sqrt(3. * gain)
      return random.uniform(rng, shape, jnp.float32, -lim, lim)
    else:
      raise ValueError('invalid distribution for ScaleInitializer')

  return Init
Returns an initializer that adjusts its scale based on weight shapes.
trax/layers/initializers.py