jax-ml/jax · error · ValueError

Expected num_classes to match the size of axis {axis}, but {

Error message

Expected num_classes to match the size of axis {axis}, but {num_classes} != {axis_size}

What it means

When jax.nn.one_hot is called under an existing pmap axis context with a string `axis` naming a pmap-mapped axis, num_classes must equal that mapped axis's size; otherwise the one-hot cannot be built from lax.axis_index. The error reports the mismatch.

Source

Thrown at jax/_src/nn/functions.py:727

    else:
      raise ValueError(
          f"Unknown algorithm '{algorithm}'. Expected 'fast' or 'stable'.")
  return jnp.subtract(x, mean) * lax.rsqrt(variance + epsilon)

# TODO(slebedev): Change the type of `x` to `ArrayLike`.
@api.jit(static_argnames=("num_classes", "dtype", "axis", "out_sharding"))
def _one_hot(x: Array, num_classes: int, *,
             dtype: DTypeLike, axis: int | AxisName,
             out_sharding: NamedSharding | None) -> Array:
  num_classes = core.concrete_dim_or_error(
      num_classes,
      "The error arose in jax.nn.one_hot argument `num_classes`.")
  try:
    out_axis = util.canonicalize_axis(axis, x.ndim + 1)  # pyrefly: ignore[bad-argument-type]
  except TypeError:
    axis_size = lax.axis_size(axis)
    if num_classes != axis_size:
      raise ValueError(f"Expected num_classes to match the size of axis {axis}, "
                       f"but {num_classes} != {axis_size}") from None
    axis_idx = lax.axis_index(axis)
    return jnp.asarray(x == axis_idx, dtype=dtype)
  assert isinstance(axis, SupportsIndex)
  axis = operator.index(axis)
  lhs = lax.expand_dims(x, (axis,))
  rhs_shape = [1] * x.ndim
  rhs_shape.insert(out_axis, num_classes)
  x_aval = core.typeof(x)
  rhs_spec = [None] * len(rhs_shape)
  if out_sharding is None:
    rhs_sharding = NamedSharding(x_aval.sharding.mesh, P(*rhs_spec))
  else:
    if out_sharding.spec.unreduced or out_sharding.spec.reduced:
      raise NotImplementedError
    out_x_spec = out_sharding.spec[:out_axis] + out_sharding.spec[out_axis+1:]
    out_x_spec = P(*out_x_spec)._normalized_spec_for_aval(x_aval.ndim)
    if out_x_spec != x_aval.sharding.spec:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Set num_classes equal to the pmap axis size (e.g. jax.local_device_count() or the size passed to pmap for that axis name)
Avoid naming the one-hot axis the same as a pmap axis if you did not intend the mapped behavior (use an int axis)

Example fix

// before
jax.pmap(lambda x: jax.nn.one_hot(x, 8, axis='i'), axis_name='i')(data)  # 4 devices

// after
n = jax.local_device_count()
jax.pmap(lambda x: jax.nn.one_hot(x, n, axis='i'), axis_name='i')(data)

Defensive patterns

Strategy: validation

Validate before calling

# before pmap'd one_hot with named axis:
axis_size = jax.local_device_count()  # or explicit pmap axis size
assert num_classes == axis_size, f'num_classes {num_classes} != pmap axis size {axis_size}'

Prevention

Never name the one-hot axis the same as a pmap axis name unless intentional
Derive num_classes from the device/axis size instead of hardcoding

When it happens

Trigger: Calling jax.nn.one_hot(x, num_classes, axis='i') inside jax.pmap(axis_name='i') where num_classes != the number of devices mapped to axis 'i'.

Common situations: Data/model parallel one-hot generation under pmap where num_classes was configured for a different device count; changing device count without updating num_classes.

Related errors

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