jax-ml/jax · error · ValueError

in_specs containing unreduced {spec} passed to shard_map sho

Error message

in_specs containing unreduced {spec} passed to shard_map should be equal to the unreduced present on the in_aval {aval.str_short(True)}

What it means

When an input array carries a sharding whose spec includes `unreduced` names, the in_specs passed to shard_map must declare exactly the same unreduced set. A mismatch means the caller is trying to reinterpret replication semantics inconsistently.

Source

Thrown at jax/_src/shard_map.py:855

    out_avals_lo = out_avals
  out = trace.emit_eqn([*const_tracers, *in_tracers], out_avals_lo, prim, params,
                       effs, source_info)
  if trace.requires_low:
    out = pe.raise_lo_outs(out_avals, out)
  return out_avals_ft.update(out)
pe.DynamicJaxprTrace.process_shard_map = _shard_map_staging

# TODO add underscore version, for direct-linearize to consume

def _spec_to_names(spec: PartitionSpec):
  return {i: names if isinstance(names, tuple) else (names,)
          for i, names in enumerate(spec.partitions) if names is not None}

def _shard_shaped_array(mesh: Mesh, manual_axes: frozenset, check_vma,
                        spec, aval: core.ShapedArray) -> core.ShapedArray:
  assert isinstance(aval, core.ShapedArray)
  if spec.unreduced != aval.sharding.spec.unreduced:
    raise ValueError(
        f"in_specs containing unreduced {spec} passed to shard_map should be"
        " equal to the unreduced present on the in_aval"
        f" {aval.str_short(True)}")
  if spec.unreduced_kind is not aval.sharding.spec.unreduced_kind:
    raise ValueError(
        f"in_specs containing unreduced_kind {spec} passed to shard_map should"
        " be equal to the unreduced_kind present on the in_aval"
        f" {aval.str_short(True)}")
  if spec.reduced != aval.sharding.spec.reduced:
    raise ValueError(
        f"in_specs containing reduced {spec} passed to shard_map should be"
        f" equal to the reduced present on the in_aval {aval.str_short(True)}")
  names = _spec_to_names(spec)
  new_shape = tuple(sz // prod(mesh.shape[n] for n in names.get(i, ()))
                    for i, sz in enumerate(aval.shape))
  manual_mesh = _as_manual_mesh(mesh, manual_axes)
  new_sharding = aval.sharding.update(
      mesh=manual_mesh,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Make in_specs' unreduced tuple equal to the input array's aval sharding unreduced tuple
  2. Re-create the input with the intended unreduced sharding before the call
  3. Inspect aval.str_short(True) / arr.sharding to see the expected unreduced names

Example fix

// before
y = shard_map(f, mesh, x_unreduced_on_r, in_specs=P('d'))
// after
y = shard_map(f, mesh, x_unreduced_on_r, in_specs=P('d', unreduced=('r',)))
Defensive patterns

Strategy: validation

Validate before calling

def unreduced_matches(x, spec):
    return getattr(getattr(x, 'sharding', None), 'spec', None) is None or spec.unreduced == x.sharding.spec.unreduced

Type guard

def in_spec_consistent(x, spec) -> bool:
    s = getattr(getattr(x, 'sharding', None), 'spec', None)
    return s is None or spec.unreduced == s.unreduced

Try / catch

try: shard_map(...) except ValueError as e: if 'unreduced' in str(e) and 'in_specs' in str(e): read aval via jax.core.shaped_abstractify and copy its unreduced into spec; else: raise

Prevention

When it happens

Trigger: Passing an array with NamedSharding spec P(..., unreduced=('r',)) into shard_map with in_specs=P(...) whose unreduced field is () or a different tuple.

Common situations: Chaining shard_map calls or jit->shard_map pipelines where the input was produced with unreduced semantics; hand-constructing specs that drift from the array's stored sharding.

Related errors

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