{"record":{"id":"a3794f0553fc73f3","repo":"jax-ml/jax","slug":"expected-num-classes-to-match-the-size-of-axis-ax","errorCode":null,"errorMessage":"Expected num_classes to match the size of axis {axis}, but {num_classes} != {axis_size}","messagePattern":"Expected num_classes to match the size of axis (.+?), but (.+?) != (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/nn/functions.py","lineNumber":727,"sourceCode":"    else:\n      raise ValueError(\n          f\"Unknown algorithm '{algorithm}'. Expected 'fast' or 'stable'.\")\n  return jnp.subtract(x, mean) * lax.rsqrt(variance + epsilon)\n\n# TODO(slebedev): Change the type of `x` to `ArrayLike`.\n@api.jit(static_argnames=(\"num_classes\", \"dtype\", \"axis\", \"out_sharding\"))\ndef _one_hot(x: Array, num_classes: int, *,\n             dtype: DTypeLike, axis: int | AxisName,\n             out_sharding: NamedSharding | None) -> Array:\n  num_classes = core.concrete_dim_or_error(\n      num_classes,\n      \"The error arose in jax.nn.one_hot argument `num_classes`.\")\n  try:\n    out_axis = util.canonicalize_axis(axis, x.ndim + 1)  # pyrefly: ignore[bad-argument-type]\n  except TypeError:\n    axis_size = lax.axis_size(axis)\n    if num_classes != axis_size:\n      raise ValueError(f\"Expected num_classes to match the size of axis {axis}, \"\n                       f\"but {num_classes} != {axis_size}\") from None\n    axis_idx = lax.axis_index(axis)\n    return jnp.asarray(x == axis_idx, dtype=dtype)\n  assert isinstance(axis, SupportsIndex)\n  axis = operator.index(axis)\n  lhs = lax.expand_dims(x, (axis,))\n  rhs_shape = [1] * x.ndim\n  rhs_shape.insert(out_axis, num_classes)\n  x_aval = core.typeof(x)\n  rhs_spec = [None] * len(rhs_shape)\n  if out_sharding is None:\n    rhs_sharding = NamedSharding(x_aval.sharding.mesh, P(*rhs_spec))\n  else:\n    if out_sharding.spec.unreduced or out_sharding.spec.reduced:\n      raise NotImplementedError\n    out_x_spec = out_sharding.spec[:out_axis] + out_sharding.spec[out_axis+1:]\n    out_x_spec = P(*out_x_spec)._normalized_spec_for_aval(x_aval.ndim)\n    if out_x_spec != x_aval.sharding.spec:","sourceCodeStart":709,"sourceCodeEnd":745,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L709-L745","documentation":"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.","triggerScenarios":"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'.","commonSituations":"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.","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)"],"exampleFix":"// before\njax.pmap(lambda x: jax.nn.one_hot(x, 8, axis='i'), axis_name='i')(data)  # 4 devices\n\n// after\nn = jax.local_device_count()\njax.pmap(lambda x: jax.nn.one_hot(x, n, axis='i'), axis_name='i')(data)","handlingStrategy":"validation","validationCode":"# before pmap'd one_hot with named axis:\naxis_size = jax.local_device_count()  # or explicit pmap axis size\nassert num_classes == axis_size, f'num_classes {num_classes} != pmap axis size {axis_size}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","nn","one-hot","pmap","distributed"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}