jax-ml/jax · error · ValueError

{name} cannot accept args which are unreduced. Got {a.str_sh

Error message

{name} cannot accept args which are unreduced. Got {a.str_short(True)} and axes={axes}

What it means

check_unreduced_args guards collective/axis-consuming primitives: arguments marked unreduced on any of the given mesh axes cannot be consumed there, because the value is only valid un-reduced on those axes. E.g. a psum-intermediate consumed again along its reduction axis is rejected.

Source

Thrown at jax/_src/core.py:2749

    return x
  cur_mesh = mesh_lib.get_abstract_mesh()
  if not config._check_vma.value and all(a in cur_mesh.manual_axes for a in axes):
    return x
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshard/convert the value so it is no longer unreduced on those axes (e.g. insert_reduced_reshard or an explicit reshard) before consuming it
2. Restructure the computation to consume the pre-reduction value or use the reduced result instead
3. For custom primitives, mark the operand axes properly so the framework inserts resharding

Example fix

// before
y = jax.lax.psum(x, 'i')   # x unreduced on 'i'
z = jax.lax.psum(y, 'i')   # consuming unreduced on same axis

// after
y = jax.lax.psum(x, 'i')
y = reshard(y)             # clear unreduced state on 'i'
z = some_local_op(y)

if any(a.mat.unreduced & axes for a in args):
    args = [reshard(a) for a in args]

def args_clean_of_unreduced(args, axes): return not any(a.mat.unreduced & set(axes) for a in args)

When it happens

Trigger: Inside shard_map/spmd code, feeding an array whose mat.unreduced intersects the axes a primitive operates over — e.g. using an unreduced psum output in another collective over the same axis without resharding.

Common situations: Chaining collectives (psum output reused in another psum/all_gather) on the same mesh axis; writing custom spmd primitives without inserting resharding.

Related errors

• {name} cannot accept args with unreduced_kind={a.mat.unreduc
• {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
• Sharding on sequence dim is not allowed.

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