jax-ml/jax · error · ValueError

cannot accept args which are reduced. Got and axes=

Error message

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

What it means

check_unreduced_args' third rule: arguments marked reduced on any of the given axes cannot be consumed by primitives operating on those axes — the data is already reduced there and reusing it as if it varied is invalid.

Solutions

1. Reshard the intermediate to clear its 'reduced' marking on those axes
2. Use the original pre-reduction tensor for follow-up collectives
3. For custom primitives, declare proper resharding effects

Example fix

// before
r = jax.lax.psum(x, 'i')
r2 = axis_op(r, axis='i')  # r is reduced on 'i'

// after
r = jax.lax.psum(x, 'i')
r = reshard(r)
r2 = axis_op(r, axis='i')

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

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

When it happens

Trigger: Using the output of a collective (marked reduced on axis 'x') as input to another primitive over 'x' inside shard_map/spmd without resharding the intermediate.

Common situations: Composing psum/all_reduce results into further axis-wise ops; custom spmd code that forwards outputs of collectives directly into more collectives.

Related errors

• cannot accept args which are unreduced. Got and axes=
• cannot accept args with unreduced_kind= . Expected…
• all_gather_reduced only accepts inputs that are varying. Got
• all_gather_reduced's input cannot be reduced across the…
• callbacks are only supported in spmd computations when all…

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

Appendix: source

Thrown at jax/_src/core.py:2758

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]
  out_reduced = frozenset.union(*in_reduced)
  out = []
  for arg, src_reduced in zip(args, in_reduced):
    aval = shaped_abstractify(arg)
    if (isinstance(aval, ShapedArray) and aval.ndim == 0 and out_reduced and

View on GitHub (pinned to 1e1c6a8fc0)