jax-ml/jax · error · ShardingTypeError

The input part of spec in out_sharding should match the spec

Error message

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}

What it means

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.

Source

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

                       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:
      raise core.ShardingTypeError(
          "The input part of spec in out_sharding should match the spec of"
          f" `x`. Got typeof(x).spec={x_aval.sharding.spec} and"
          f" out_sharding.spec={out_x_spec}")
    rhs_spec[out_axis] = out_sharding.spec[out_axis]
    rhs_sharding = out_sharding.update(spec=P(*rhs_spec))
  rhs = lax.broadcasted_iota(x.dtype, rhs_shape, out_axis, out_sharding=rhs_sharding)
  return (lhs == rhs).astype(dtype)

# TODO(slebedev): Change the type of `x` to `ArrayLike`.
def one_hot(x: Any, num_classes: int, *,
            dtype: Any | None = None, axis: int | AxisName = -1,
            out_sharding: NamedSharding | P | None = None) -> Array:
  r"""One-hot encodes the given indices.

  Each index in the input ``x`` is encoded as a vector of zeros of length
  ``num_classes`` with the element at ``index`` set to one::

    >>> jax.nn.one_hot(jnp.array([0, 1, 2]), 3)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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)))
2. Construct out_sharding by taking x's spec and inserting the desired spec entry at the output axis position
3. Ensure out_sharding.spec has no unreduced/reduced axes (that path raises NotImplementedError first)

Example fix

// before
out = jax.nn.one_hot(x, 10, axis=1, out_sharding=NamedSharding(mesh, P('x', 'y')))  # x is P('x', None)

// after
x = jax.device_put(x, NamedSharding(mesh, P('x', None)))
out = jax.nn.one_hot(x, 10, axis=1, out_sharding=NamedSharding(mesh, P('x', 'y', None)))

from jax.sharding import NamedSharding, PartitionSpec as P

def make_out_sharding(x_aval, mesh, new_axis, new_axis_spec):
    spec = list(x_aval.sharding.spec)
    spec.insert(new_axis, new_axis_spec)
    return NamedSharding(mesh, P(*spec))

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• Sharding spec {spec} implies that array axis {dim} is partit
• {}: arrays must have same number of dimensions, got {}.
• out_sharding passed to {name} cannot be reduced. Got {out_sh
• Reduced axes can only refer to mesh axes that is of type `Ex
• Expected num_classes to match the size of axis {axis}, but {

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