jax-ml/jax · error · ValueError

in_specs containing unreduced_kind {spec} passed to shard_ma

Error message

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

What it means

Input arrays with unreduced sharding also carry an `unreduced_kind` (e.g. how the replicated partial values combine). shard_map requires the in_specs' unreduced_kind to equal the aval's stored unreduced_kind; a mismatch raises this ValueError.

Source

Thrown at jax/_src/shard_map.py:860

  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,
      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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Align unreduced_kind in in_specs with the input's sharding (inspect arr.sharding.spec.unreduced_kind)
2. Rebuild the input array's sharding with the desired unreduced_kind before shard_map
3. Avoid hand-building unreduced shardings; use the same helper on both sides

Example fix

// before
y = shard_map(f, mesh, x, in_specs=P(unreduced=('r',), unreduced_kind='mul'))  # x is 'add'
// after
y = shard_map(f, mesh, x, in_specs=P(unreduced=('r',), unreduced_kind='add'))

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

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

When it happens

Trigger: Passing an array whose sharding was created with a different unreduced_kind (e.g. 'add' vs 'mul') than the unreduced_kind declared in in_specs.

Common situations: Mixing APIs that default unreduced_kind differently across JAX versions; constructing NamedSharding manually with mismatched kind parameters.

Related errors

• out_specs passed to shard_map should be equal to the unreduc
• in_specs containing unreduced {spec} passed to shard_map sho
• in_specs containing reduced {spec} passed to shard_map shoul
• out_specs passed to shard_map should be equal to the unreduc
• out_specs passed to shard_map should be equal to the reduced

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