code: def variance_scaling(
      scale: RealNumeric,
      mode: Literal["fan_in"] | Literal["fan_out"] | Literal["fan_avg"],
      distribution: (
          Literal["truncated_normal"]
          | Literal["normal"]
          | Literal["uniform"]
      ),
      in_axis: int | Sequence[int] = -2,
      out_axis: int | Sequence[int] = -1,
      batch_axis: Sequence[int] = (),
      dtype: DTypeLikeInexact = jnp.float_,
  ) -> Initializer:
      r"""
      Initializer that adapts its scale to the shape of the weights tensor.
  
      With ``distribution="truncated_normal"`` or ``distribution="normal"``, samples
      are drawn from a (truncated) normal distribution with a mean of zero
      and a standard deviation (after truncation, if applicable) of
      :math:`\sqrt{\frac{scale}{n}}`, where `n` is:
  
      * the number of input units in the weights tensor, if ``mode="fan_in"``,
      * the number of output units, if ``mode="fan_out"``, or
      * the average of the numbers of input and output units, if ``mode="fan_avg"``.
  
      This initializer can be configured with ``in_axis``, ``out_axis``, and
      ``batch_axis`` to work with general convolutional or dense layers; axes that
      are not in any of those arguments are assumed to be the "receptive field"
      (convolution kernel spatial axes).
  
      With ``distribution="truncated_normal"``, the absolute values of the samples
      are truncated at 2 standard deviations before scaling.
  
      With ``distribution="uniform"``, samples are drawn from:
  
      * a uniform interval, if `dtype` is real, or
      * a uniform disk, if `dtype` is complex,
  
      with a mean of zero and a standard deviation of :math:`\sqrt{\frac{scale}{n}}`
      where `n` is defined above.
  
      Args:
        scale: scaling factor (positive float).
        mode: one of ``"fan_in"``, ``"fan_out"``, and ``"fan_avg"``.
        distribution: random distribution to use. One of ``"truncated_normal"``,
          ``"normal"`` and ``"uniform"``.
        in_axis: axis or sequence of axes of the input dimension in the weights
          array.
        out_axis: axis or sequence of axes of the output dimension in the weights
          array.
        batch_axis: axis or sequence of axes in the weight array that should be
          ignored.
        dtype: the dtype of the weights.
      """
  
      def init(
          key: KeyArray,
          shape: core.Shape,
          dtype: DTypeLikeInexact = dtype,
      ) -> Array:
          dtype = dtypes.canonicalize_dtype(dtype)
          named_shape = core.as_named_shape(shape)
          fan_in, fan_out = _compute_fans(
              named_shape, in_axis, out_axis, batch_axis
          )
          if mode == "fan_in":
              denominator = fan_in
          elif mode == "fan_out":
              denominator = fan_out
          elif mode == "fan_avg":
              denominator = (fan_in + fan_out) / 2
          else:
              raise ValueError(
                  f"invalid mode for variance scaling initializer: {mode}"
              )
          variance = jnp.array(scale / denominator, dtype=dtype)
  
          if distribution == "truncated_normal":
              if jnp.issubdtype(dtype, jnp.floating):
                  # constant is stddev of standard normal truncated to (-2, 2)
                  stddev = jnp.sqrt(variance) / jnp.array(
                      0.87962566103423978, dtype
                  )
                  return (
                      random.truncated_normal(
                          key, -2, 2, named_shape, dtype
                      )
                      * stddev
                  )
              else:
                  # constant is stddev of complex standard normal truncated to 2
                  stddev = jnp.sqrt(variance) / jnp.array(
                      0.95311164380491208, dtype
                  )
                  return (
                      _complex_truncated_normal(
                          key, 2, named_shape, dtype
                      )
                      * stddev
                  )
          elif distribution == "normal":
              return random.normal(key, named_shape, dtype) * jnp.sqrt(
                  variance
              )
          elif distribution == "uniform":
              if jnp.issubdtype(dtype, jnp.floating):
                  return random.uniform(
                      key, named_shape, dtype, -1
                  ) * jnp.sqrt(3 * variance)
              else:
                  return _complex_uniform(
                      key, named_shape, dtype
                  ) * jnp.sqrt(variance)
          else:
              raise ValueError(
                  f"invalid distribution for variance scaling initializer: {distribution}"
              )
  
      return init
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import math
  import typing
  from typing import Any, Literal, Protocol
  import numpy
  import jax.numpy
  from jax import lax
  from jax import random
  from jax._src import core
  from jax._src import dtypes
  from jax._src.typing import Array, ArrayLike
  from jax._src.util import set_module
label_desc: Use image node_id 10 for calling a global function with example usage: variance_scaling(scale, mode, distribution, in_axis, out_axis, batch_axis, dtype) and returns: init, unknown, unknown, unknown, unknown, unknown