jax-ml/jax · error · ValueError

invalid distribution for variance scaling initializer: {dist

Error message

invalid distribution for variance scaling initializer: {distribution}

What it means

variance_scaling's second enum argument, `distribution`, supports 'truncated_normal', 'normal', and 'uniform'. Any other string raises this ValueError at init call time.

Source

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

        # 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):
        return random.uniform(key, shape, dtype, -1,
                              out_sharding=out_sharding) * jnp.sqrt(3 * variance)
      else:
        return _complex_uniform(key, shape, dtype) * jnp.sqrt(variance)
    else:
      raise ValueError(f"invalid distribution for variance scaling initializer: {distribution}")

  return init

@export
def glorot_uniform(in_axis: int | Sequence[int] = -2,
                   out_axis: int | Sequence[int] = -1,
                   batch_axis: int | Sequence[int] = (),
                   dtype: DTypeLikeInexact | None = None) -> Initializer:
  """Builds a Glorot uniform initializer (aka Xavier uniform initializer).

  A `Glorot uniform initializer`_ is a specialization of
  :func:`jax.nn.initializers.variance_scaling` where ``scale = 1.0``,
  ``mode="fan_avg"``, and ``distribution="uniform"``.

  Args:
    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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use exactly 'truncated_normal', 'normal', or 'uniform'
If you need a different sampler, wrap the initializer and post-transform the samples

Example fix

// before
init = jax.nn.initializers.variance_scaling(1.0, 'fan_in', 'trunc_normal')

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

Defensive patterns

Strategy: validation

Validate before calling

DISTS = ('truncated_normal', 'normal', 'uniform')
assert distribution in DISTS, f'distribution must be one of {DISTS}'

Type guard

def is_valid_distribution(d: str) -> bool: return d in ('truncated_normal','normal','uniform')

Prevention

Use Literal typing for initializer enum args
Validate initializer configs once at parse time

When it happens

Trigger: Passing distribution='tn', 'gaussian', distribution=None, or 'truncated-normal' (hyphen typo).

Common situations: Config-driven initializer construction with typo'd strings; porting from libraries that use 'trunc_normal' naming.

Related errors

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