jax-ml/jax · error · ValueError

{name} cannot accept args with unreduced_kind={a.mat.unreduc

Error message

{name} cannot accept args with unreduced_kind={a.mat.unreduced_kind}. Expected unreduced_kind={kind}

What it means

check_unreduced_args also enforces that all unreduced operands share the same UnreducedKind as the primitive expects (default sum). An operand marked unreduced with kind=min cannot feed a primitive expecting sum-style unreduce semantics.

Source

Thrown at jax/_src/core.py:2753

  new_axes = axes if cur_mesh.empty else order_wrt_mesh(cur_mesh, axes)
  assert set(new_axes) == set(axes)
  del axes
  return tree_map(lambda leaf: reduced_vary_cast_p.bind(leaf, axes=new_axes), x)

reduced_vary_cast_p = Primitive('reduced_vary_cast_p')

#######################################################################

def check_unreduced_args(args, axes, name, kind=UnreducedKind.sum):
  axes = axes if isinstance(axes, (tuple, list)) else (axes,)
  axes = set(axes)
  for a in args:
    if a.mat.unreduced & axes:
      raise ValueError(
          f"{name} cannot accept args which are unreduced. Got"
          f" {a.str_short(True)} and axes={axes}")
    if a.mat.unreduced and a.mat.unreduced_kind is not kind:
      raise ValueError(
          f"{name} cannot accept args with"
          f" unreduced_kind={a.mat.unreduced_kind}. Expected"
          f" unreduced_kind={kind}")
    if a.mat.reduced & axes:
      raise ValueError(
          f"{name} cannot accept args which are reduced. Got"
          f" {a.str_short(True)} and axes={axes}")

def insert_reduced_reshard(args):
  cur_mesh = mesh_lib.get_abstract_mesh()
  if not cur_mesh.are_all_axes_explicit:
    return args
  # TODO(yashkatariya): Handle >2 args too
  if len(args) != 2:
    return args
  in_reduced = [aval.sharding.spec.reduced
                if isinstance(aval := shaped_abstractify(a), ShapedArray)
                else frozenset() for a in args]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Insert a reshard so operands lose their unreduced marking before the next collective
2. Match the primitive's expected kind by transforming the value appropriately (e.g. re-reduce with the same kind)
3. Pass the correct kind parameter to the primitive if it accepts one

Example fix

// before
m = jax.lax.pmin(x, 'i')       # unreduced_kind=min
s = consume_expecting_sum(m)   # rejected

// after
m = jax.lax.pmin(x, 'i')
m = reshard(m)
s = consume_expecting_sum(m)

if any(a.mat.unreduced and a.mat.unreduced_kind is not kind for a in args):
    args = [reshard(a) for a in args]

def kinds_match(args, kind): return all(a.mat.unreduced_kind is kind for a in args if a.mat.unreduced)

When it happens

Trigger: Mixing pmin/pmax-produced unreduced values into primitives expecting UnreducedKind.sum (or vice versa) inside shard_map/spmd pipelines.

Common situations: Chaining different reduction flavors (pmin then psum paths) without resharding between; custom collectives defaulting to sum while receiving min-kind operands.

Related errors

• {name} cannot accept args which are unreduced. Got {a.str_sh
• {name} cannot accept args which are reduced. Got {a.str_shor
• callbacks are only supported in spmd computations when all m
• callbacks do not support specifying sharding inside spmd com
• Expected unreduced_kind to be of type `jax.sharding.Unreduce

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