jax-ml/jax · error · ValueError

reduced_vary_cast input cannot be varying across the axis_na

Error message

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

What it means

`reduced_vary_cast` requires the input not be varying across the axes being cast; the abstract eval rejects inputs where `aval.mat.varying & set(axes)` is non-empty, because casting a varying axis to varying is contradictory.

Source

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

    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)
  return aval.update(manual_axis_type=aval.mat.update(
    varying=out_vma, reduced=new_reduced))
core.reduced_vary_cast_p.def_abstract_eval(_reduced_vary_cast_abstract_eval)

def _reduced_vary_cast_transpose_rule(cts, x, *, axes):
  assert ad.is_undefined_primal(x)
  return (vary_unreduced_cast(cts, axis_name=axes),)
ad.deflinear2(core.reduced_vary_cast_p, _reduced_vary_cast_transpose_rule)

def _reduced_vary_cast_batcher(vals_in, dims_in, *, axes):
  raise NotImplementedError

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. First reduce over that axis (e.g. psum) or use the appropriate Varying->Varying/Unreduced cast
2. Verify `jax.typeof(x).mat` at each stage and fix the preceding op
3. Use `pcast(..., to='varying')` for state-aware dispatch

Example fix

// before
y = reduced_vary_cast(x, 'dev')  # x varying over 'dev'
// after
y = vary_appropriate_op(x, 'dev')  # reduce or leave varying; pcast(x,'dev',to='varying') is a no-op path

t = jax.typeof(x)
assert not (set(axes) & t.mat.varying), 'input varies over cast axes; reduce first'

def not_varying_over(x, axes) -> bool:
    return not (set(axes) & jax.typeof(x).mat.varying)

When it happens

Trigger: Calling reduced_vary_cast over an axis the input still varies along (e.g. after gathering/splitting re-introduced variation, or on a fresh broadcasted value).

Common situations: State-machine mistakes in manual axis tracking where a value is believed reduced but a preceding op actually re-varied it.

Related errors

• vary_unreduced_cast only accepts inputs that are varying. Go
• reduced_vary_cast only accepts inputs that are reduced. Got
• 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/69d4dde8fc874ec1. Report an issue: GitHub.