jax-ml/jax · error · ValueError

{name} cannot accept args which are reduced. Got {a.str_shor

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.

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)

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

• {name} cannot accept args which are unreduced. Got {a.str_sh
• {name} cannot accept args with unreduced_kind={a.mat.unreduc
• 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/4bea3a1bf779c257. Report an issue: GitHub.