jax-ml/jax · error · ValueError

vary_unreduced_cast input cannot be unreduced across the axi

Error message

vary_unreduced_cast input cannot be unreduced across the axis_name provided. Got x={aval.str_short(True)} and axis_name={axes}

What it means

`vary_unreduced_cast` cannot mark an axis unreduced if the input is already unreduced across that axis. The abstract eval rejects inputs where `aval.mat.unreduced & set(axes)` is non-empty, since the cast would be redundant/contradictory.

Source

Thrown at jax/_src/lax/parallel.py:2940

mlir.register_lowering(vary_unreduced_cast_p, lambda ctx, x, *, axes: [x])

def _vary_unreduced_cast_abstract_eval(aval, *, axes):
  assert isinstance(axes, tuple)
  _check_axis_names(axes, 'vary_unreduced_cast')
  check_unreduced_args([aval], axes, 'vary_unreduced_cast')
  if not aval.mat.varying:
    raise ValueError('vary_unreduced_cast only accepts inputs that are'
                     f' varying. Got {aval.str_short(True)}')
  # If the intersection between aval.mat.varying and axes is empty, error
  if not (aval.mat.varying & set(axes)):
    raise ValueError(
        "vary_unreduced_cast is a Varying->Unreduced collective. This"
        " means that the axis names mentioned in `axes` passed to"
        " `vary_unreduced_cast` must be present in"
        f" `jax.typeof(x).mat.varying`. Got axes={axes} and"
        f" jax.typeof(x).mat.varying={aval.mat.varying}")
  if aval.mat.unreduced & set(axes):
    raise ValueError(
        "vary_unreduced_cast input cannot be unreduced across the axis_name"
        f" provided. Got x={aval.str_short(True)} and axis_name={axes}")

  new_unreduced = aval.mat.unreduced | frozenset(axes)
  out_vma = frozenset(i for i in aval.mat.varying if i not in axes)
  return aval.update(manual_axis_type=aval.mat.update(
    varying=out_vma, unreduced=new_unreduced))
vary_unreduced_cast_p.def_abstract_eval(_vary_unreduced_cast_abstract_eval)

def _vary_unreduced_cast_transpose_rule(cts, x, *, axes):
  assert ad.is_undefined_primal(x)
  return (core.reduced_vary_cast(cts, axis_name=axes),)
ad.deflinear2(vary_unreduced_cast_p, _vary_unreduced_cast_transpose_rule)

def _vary_unreduced_cast_batcher(vals_in, dims_in, *, axes):
  raise NotImplementedError
batching.primitive_batchers[vary_unreduced_cast_p] = _vary_unreduced_cast_batcher

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the duplicate cast — the value is already unreduced across that axis
2. Guard the cast with a check on `jax.typeof(x).mat.unreduced`
3. Use `pcast`, which is idempotent-aware dispatch, instead of the raw cast

Example fix

// before
x = vary_unreduced_cast(x, 'dev')
x = vary_unreduced_cast(x, 'dev')  # duplicate
// after
x = vary_unreduced_cast(x, 'dev')

if set(axes) & jax.typeof(x).mat.unreduced:
    return x  # already unreduced; skip cast

def already_unreduced(x, axes) -> bool:
    return bool(set(axes) & jax.typeof(x).mat.unreduced)

When it happens

Trigger: Applying vary_unreduced_cast to the output of a previous vary_unreduced_cast (or unreduced_psum) over the same axis_name; double-casting in composed collective helpers.

Common situations: Wrapper functions that defensively cast to unreduced being applied to already-cast values; layered abstractions over manual collectives stacking casts.

Related errors

• vary_unreduced_cast only accepts inputs that are varying. Go
• vary_unreduced_cast is a Varying->Unreduced collective. This
• reduced_vary_cast only accepts inputs that are reduced. Got
• reduced_vary_cast is a Reduced->Varying collective. This mea
• reduced_vary_cast input cannot be varying across the axis_na

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