jax-ml/jax · error · ValueError

reduced_vary_cast only accepts inputs that are reduced. Got

Error message

reduced_vary_cast only accepts inputs that are reduced. Got {aval.str_short(True)}

What it means

`reduced_vary_cast` is a Reduced->Varying no-op cast; its abstract eval requires the input to currently be in the reduced state (`aval.mat.reduced` non-empty). Passing a value that is varying/unreduced/invarying raises this ValueError.

Source

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

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

####################### reduced_vary_cast #############################

# Reduced -> Varying no-op cast
# Traceable defined in core.py to avoid circular imports
core.reduced_vary_cast_p.def_impl(
    partial(_raise_valueerror, 'reduced_vary_cast'))
mlir.register_lowering(core.reduced_vary_cast_p, lambda ctx, x, *, axes: [x])

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

  new_reduced = frozenset(i for i in aval.mat.reduced if i not in axes)
  out_vma = aval.mat.varying | frozenset(axes)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect `jax.typeof(x).mat.reduced` and only cast values that are reduced
2. Use `jax.lax.pcast(x, axis_name, to='varying')` which picks the right cast for the actual state
3. Apply the intended reduction (psum etc.) before the cast

Example fix

// before
y = reduced_vary_cast(x, 'dev')  # x not reduced
// after
y = pcast(psum_like(x, 'dev'), 'dev', to='varying')
# or simply pcast(x, 'dev', to='varying')

if not jax.typeof(x).mat.reduced:
    x = apply_reduction(x, axes)  # real collective first

def is_reduced(x) -> bool:
    return bool(jax.typeof(x).mat.reduced)

catch ValueError, fall back to pcast(x, axes, to='varying')

When it happens

Trigger: Calling jax.lax.core reduced_vary_cast (exposed via pcast dispatch) on a value produced without a reducing collective (e.g. a plain array or an unreduced output).

Common situations: Hand-managing cast chains between axis states and calling the wrong directional cast; assuming a psum was applied when it was optimized away.

Related errors

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

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