jax-ml/jax · error · ValueError

in_specs containing reduced {spec} passed to shard_map shoul

Error message

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

What it means

When an input array's sharding spec includes `reduced` names, shard_map requires the in_specs' reduced tuple to exactly equal the aval's reduced tuple, so reduction semantics are not silently redefined at the boundary.

Source

Thrown at jax/_src/shard_map.py:865

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,
      spec=core.modify_spec_for_auto_manual(aval.sharding.spec, manual_mesh))
  vma = (_spec_to_vma(spec) if check_vma else frozenset()) | aval.mat.varying
  unreduced = aval.sharding.spec.unreduced if check_vma else frozenset()
  reduced = aval.sharding.spec.reduced if check_vma else frozenset()
  u_kind = aval.sharding.spec.unreduced_kind if check_vma else None
  mat = core.ManualAxisType(varying=vma, unreduced=unreduced, reduced=reduced,
                            unreduced_kind=u_kind)
  return aval.update(shape=new_shape, sharding=new_sharding,
                     manual_axis_type=mat)
core.shard_aval_handlers[core.ShapedArray] = _shard_shaped_array

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Copy the reduced tuple from the input's sharding spec into in_specs
2. Normalize the input (e.g. fully reduce / reshard) before entering shard_map
3. Print aval.str_short(True) to confirm what reduced names are expected

Example fix

// before
y = shard_map(f, mesh, x_reduced_y, in_specs=P('d'))
// after
y = shard_map(f, mesh, x_reduced_y, in_specs=P('d', reduced=('y',)))

def reduced_matches(x, spec):
    s = getattr(getattr(x, 'sharding', None), 'spec', None)
    return s is None or spec.reduced == s.reduced

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

When it happens

Trigger: Input array sharded with reduced=('y',) but in_specs declares reduced=() or a different axis tuple.

Common situations: Pipelines where a previous stage produced partially-reduced arrays (e.g. SP-style gradients) that are fed into another shard_map with fresh specs.

Related errors

• out_specs passed to shard_map should be equal to the reduced
• reduced in {prefix}_specs {s} can only be used when the mesh
• in_specs containing unreduced {spec} passed to shard_map sho
• in_specs containing unreduced_kind {spec} passed to shard_ma
• out_specs passed to shard_map should be equal to the unreduc

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