(out_dim, in_dim, scale, mode, distribution)
return Init
Returns an initializer that adjusts its scale based on weight shapes.
Returns an initializer that adjusts its scale based on weight shapes.
Returns an initializer that adjusts its scale based on weight shapes .
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
def ScaledInitializer ( out_dim , in_dim , scale , mode , distribution ) : 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