code: def ScaledInitializer(out_dim, in_dim, scale, mode, distribution):
      """Returns an initializer that adjusts its scale based on weight shapes."""
      if scale <= 0.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) / 0.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.0 * gain)
              return random.uniform(rng, shape, jnp.float32, -lim, lim)
          else:
              raise ValueError(
                  "invalid distribution for ScaleInitializer"
              )
  
      return Init
  
file_path: trax/trax/layers/initializers.py
imports: from absl import logging
  import numpy
  import tensorflow.compat.v2
  from trax.fastmath import numpy
  from trax.fastmath import random
label_desc: Use image node_id 6 for calling a global function with example usage: ScaledInitializer(out_dim, in_dim, scale, mode, distribution) and returns: Init, new_weights, new_weights, random