jax-ml/jax · error · ValueError

in_specs containing reduced

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.

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',)))
Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Prevention

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

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

Appendix: 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)