jax-ml/jax · error · ValueError

invalid mode for variance scaling initializer: {mode}

Error message

invalid mode for variance scaling initializer: {mode}

What it means

jax.nn.initializers.variance_scaling selects the denominator via the `mode` string: 'fan_in', 'fan_out', 'fan_avg', 'fan_geo_avg'. Any other string raises this ValueError.

Source

Thrown at jax/_src/nn/initializers.py:346

    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: Array,
           shape: core.Shape,
           dtype: DTypeLikeInexact | None = dtype,
           out_sharding: OutShardingType = None) -> Array:
    shape = core.canonicalize_shape(shape)
    dtype = dtypes.default_float_dtype() if dtype is None else dtype
    fan_in, fan_out = _compute_fans(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
    elif mode == "fan_geo_avg": denominator = (fan_in * fan_out) ** 0.5
    else:
      raise ValueError(
        f"invalid mode for variance scaling initializer: {mode}")
    variance = jnp.array(scale / denominator, dtype=dtype)

    if distribution == "truncated_normal":
      if dtypes.issubdtype(dtype, np.floating):
        # constant is stddev of standard normal truncated to (-2, 2)
        stddev = jnp.sqrt(variance) / jnp.array(.87962566103423978, dtype)
        return random.truncated_normal(key, -2, 2, shape, dtype,
                                       out_sharding=out_sharding) * stddev
      else:
        # constant is stddev of complex standard normal truncated to 2
        stddev = jnp.sqrt(variance) / jnp.array(.95311164380491208, dtype)
        return _complex_truncated_normal(key, 2, shape, dtype) * stddev
    elif distribution == "normal":
      return random.normal(key, shape, dtype,
                           out_sharding=out_sharding) * jnp.sqrt(variance)
    elif distribution == "uniform":
      if dtypes.issubdtype(dtype, np.floating):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use exactly one of: 'fan_in', 'fan_out', 'fan_avg', 'fan_geo_avg' (lowercase)
2. For Kaiming-style init use 'fan_in'; for Xavier-style balance use 'fan_avg' or 'fan_geo_avg'

Example fix

// before
init = jax.nn.initializers.variance_scaling(2.0, 'fan-in', 'truncated_normal')

// after
init = jax.nn.initializers.variance_scaling(2.0, 'fan_in', 'truncated_normal')

MODES = ('fan_in', 'fan_out', 'fan_avg', 'fan_geo_avg')
assert mode in MODES, f'mode must be one of {MODES}'

def is_valid_mode(m: str) -> bool: return m in ('fan_in','fan_out','fan_avg','fan_geo_avg')

When it happens

Trigger: Calling variance_scaling(scale, mode='FAN_IN') (case-sensitive), mode='avg', mode=None, or a typo.

Common situations: Porting configs from TF/PyTorch where mode names or casing differ; hand-written config strings from YAML/JSON typos.

Related errors

• invalid distribution for variance scaling initializer: {dist
• Unknown algorithm '{algorithm}'. Expected 'fast' or 'stable'
• Unsupported implementation option: {implementation}
• Unsupported mode: {mode}
• Can't compute input and output sizes of a {len(shape)}-dimen

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/5b75c194588e49be. Report an issue: GitHub.