jax-ml/jax · error · ValueError

Can't compute input and output sizes of a {len(shape)}-dimen

Error message

Can't compute input and output sizes of a {len(shape)}-dimensional weights tensor with default in_axis. Must be at least 2D or specify in_axis explicitly.

What it means

Variance-scaling initializers (variance_scaling, glorot_*, lecun_*) call _compute_fans to derive fan_in/fan_out. With the default in_axis=-2 they need at least a 2D shape; a 0D or 1D shape cannot yield a fan, so this ValueError is raised.

Source

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

    return random.truncated_normal(
        key, lower, upper, shape, dtype,
        out_sharding=out_sharding) * jnp.array(stddev, dtype)
  return init

@export
def _compute_fans(shape: Sequence[int],
                  in_axis: int | Sequence[int] = -2,
                  out_axis: int | Sequence[int] = -1,
                  batch_axis: int | Sequence[int] = ()
                  ) -> tuple[float, float]:
  """
  Compute effective input and output sizes for a linear or convolutional layer.

  Axes not in in_axis, out_axis, or batch_axis are assumed to constitute the
  "receptive field" of a convolution (kernel spatial dimensions).
  """
  if isinstance(in_axis, int) and in_axis == -2 and len(shape) <= 1:
    raise ValueError(
        f"Can't compute input and output sizes of a {len(shape)}-dimensional"
        " weights tensor with default in_axis. Must be at least 2D or specify"
        " in_axis explicitly."
    )

  if isinstance(in_axis, int):
    in_size = shape[in_axis]
  else:
    in_size = math.prod([shape[i] for i in in_axis])
  if isinstance(out_axis, int):
    out_size = shape[out_axis]
  else:
    out_size = math.prod([shape[i] for i in out_axis])
  if isinstance(batch_axis, int):
    batch_size = shape[batch_axis]
  else:
    batch_size = math.prod([shape[i] for i in batch_axis])
  receptive_field_size = math.prod(shape) / in_size / out_size / batch_size

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use at least a 2D shape, e.g. (n, 1) or the true (fan_in, fan_out) matrix shape
2. For 1D parameters use delta_initializer, zeros, or normal/uniform initializers instead
3. Alternatively pass in_axis explicitly (e.g. in_axis=0) to disambiguate the fan computation

Example fix

// before
w = jax.nn.initializers.lecun_normal()(key, (256,))

// after
w = jax.nn.initializers.lecun_normal()(key, (256, 1))
# or for vectors:
b = jax.nn.initializers.normal()(key, (256,))

def fan_init(init_fn, key, shape):
    if len(shape) < 2:
        raise ValueError('fan-based initializers need >=2D shape')
    return init_fn(key, shape)

def is_fan_shape(shape) -> bool: return len(shape) >= 2

When it happens

Trigger: Calling jax.nn.initializers.glorot_uniform()(key, (n,)) or variance_scaling(...)(key, (3,)) — a single-dimension shape — with default in_axis.

Common situations: Initializing 1D bias-like vectors or scalar parameters with a fan-based initializer; passing a flattened shape accidentally; layer code where the weight shape is computed wrong.

Related errors

• orthogonal initializer requires at least a 2D shape
• {name} ndim should be {len(shape)}, but got {t.ndim}
• {name} shape should be {shape}: but got {t.shape}
• The number of query heads must be a multiple of key/value he
• scaled_matmul requires all inputs to be 3-dimensional arrays

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