{"record":{"id":"2a9bddf3da463225","repo":"jax-ml/jax","slug":"the-input-part-of-spec-in-out-sharding-should-matc","errorCode":null,"errorMessage":"The input part of spec in out_sharding should match the spec of `x`. Got typeof(x).spec={x_aval.sharding.spec} and out_sharding.spec={out_x_spec}","messagePattern":"The input part of spec in out_sharding should match the spec of `x`\\. Got typeof\\(x\\)\\.spec=(.+?) and out_sharding\\.spec=(.+?)","errorType":"validation","errorClass":"ShardingTypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/nn/functions.py","lineNumber":746,"sourceCode":"                       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:\n      raise core.ShardingTypeError(\n          \"The input part of spec in out_sharding should match the spec of\"\n          f\" `x`. Got typeof(x).spec={x_aval.sharding.spec} and\"\n          f\" out_sharding.spec={out_x_spec}\")\n    rhs_spec[out_axis] = out_sharding.spec[out_axis]\n    rhs_sharding = out_sharding.update(spec=P(*rhs_spec))\n  rhs = lax.broadcasted_iota(x.dtype, rhs_shape, out_axis, out_sharding=rhs_sharding)\n  return (lhs == rhs).astype(dtype)\n\n# TODO(slebedev): Change the type of `x` to `ArrayLike`.\ndef one_hot(x: Any, num_classes: int, *,\n            dtype: Any | None = None, axis: int | AxisName = -1,\n            out_sharding: NamedSharding | P | None = None) -> Array:\n  r\"\"\"One-hot encodes the given indices.\n\n  Each index in the input ``x`` is encoded as a vector of zeros of length\n  ``num_classes`` with the element at ``index`` set to one::\n\n    >>> jax.nn.one_hot(jnp.array([0, 1, 2]), 3)","sourceCodeStart":728,"sourceCodeEnd":764,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/nn/functions.py#L728-L764","documentation":"When jax.nn.one_hot is given an out_sharding, the part of its PartitionSpec that describes the input `x` (spec with the new one-hot axis removed and normalized) must exactly match the sharding spec of x. If they differ, this ShardingTypeError is raised because the output layout is incompatible with where x lives.","triggerScenarios":"Calling jax.nn.one_hot(x, ..., out_sharding=NamedSharding(mesh, P(...))) where the out_sharding spec, after dropping the inserted axis, does not equal typeof(x).sharding.spec (e.g. x is replicated but out_sharding shards its dims, or axis names/order differ).","commonSituations":"Specifying out_sharding on multi-host pipelines without first constraining/sharding x consistently; inserting the one-hot axis at a different position than assumed; spec mismatch after update() or P normalization.","solutions":["Shard x with the same spec (sans the new axis) before calling one_hot, e.g. x = jax.device_put(x, NamedSharding(mesh, P(*spec_without_new_axis)))","Construct out_sharding by taking x's spec and inserting the desired spec entry at the output axis position","Ensure out_sharding.spec has no unreduced/reduced axes (that path raises NotImplementedError first)"],"exampleFix":"// before\nout = jax.nn.one_hot(x, 10, axis=1, out_sharding=NamedSharding(mesh, P('x', 'y')))  # x is P('x', None)\n\n// after\nx = jax.device_put(x, NamedSharding(mesh, P('x', None)))\nout = jax.nn.one_hot(x, 10, axis=1, out_sharding=NamedSharding(mesh, P('x', 'y', None)))","handlingStrategy":"validation","validationCode":"from jax.sharding import NamedSharding, PartitionSpec as P\n\ndef make_out_sharding(x_aval, mesh, new_axis, new_axis_spec):\n    spec = list(x_aval.sharding.spec)\n    spec.insert(new_axis, new_axis_spec)\n    return NamedSharding(mesh, P(*spec))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always derive out_sharding from typeof(x).spec by inserting the new axis","Device_put x to the matching NamedSharding before calling one_hot with out_sharding"],"tags":["jax","nn","one-hot","sharding","named-sharding"],"backgroundTag":"sharding-spec-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}