jax-ml/jax · error · ValueError

{name} can only accept axis_name which corresponds to one of

Error message

{name} can only accept axis_name which corresponds to one of varying, unreduced, reduced or invarying state of the input. Got input type: {aval}, axes: {axes} and input state: {out}

What it means

`jax.lax.pcast` inspects, for each axis in `axes`, which axis state (varying/unreduced/reduced/invarying) the input is in via `_get_from`. If the requested axes map to more than one distinct state, the cast source is ambiguous and pcast raises this ValueError.

Source

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

  raise NotImplementedError
batching.primitive_batchers[core.reduced_vary_cast_p] = _reduced_vary_cast_batcher

################################## pcast #############################

def _get_from(aval, axes: tuple[AxisName, ...], name) -> str:
  out = set()
  for a in axes:
    if a in aval.mat.varying:
      out.add('varying')
    elif a in aval.mat.unreduced:
      out.add('unreduced')
    elif a in aval.mat.reduced:
      out.add('reduced')
    else:
      out.add('invarying')

  if len(out) > 1:
    raise ValueError(
        f"{name} can only accept axis_name which corresponds to one of"
        " varying, unreduced, reduced or invarying state of the input. Got"
        f" input type: {aval}, axes: {axes} and input state: {out}")
  o, = out
  return o


_pcast_funcs = {
    ('invarying', 'varying'): core.pvary,
    ('invarying', 'reduced'): preduced,
    ('varying', 'unreduced'): vary_unreduced_cast,
    ('reduced', 'varying'): core.reduced_vary_cast,
}

_allowed_pcast_to = {'unreduced', 'reduced', 'varying'}

def pcast(x, axis_name, *, to: str):
  if isinstance(axis_name, (set, frozenset)):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split the pcast into one call per axis (or per group of axes sharing the same state)
2. First normalize the input to a single state using the explicit casts (vary_unreduced_cast/reduced_vary_cast/...) per axis
3. Inspect `jax.typeof(x).mat` and re-order preceding collectives so all requested axes share one state

Example fix

# before
y = jax.lax.pcast(x, ('i', 'j'), to='unreduced')
# after
y = jax.lax.pcast(jax.lax.pcast(x, 'i', to='unreduced'), 'j', to='unreduced')

states = {jax.typeof(x).mat.state_of(a) for a in axes}  # conceptual
# practically: check each axis's membership in varying/unreduced/reduced sets and group

def axes_share_state(x, axes) -> bool:
    m = jax.typeof(x).mat
    sigs = {(a in m.varying, a in m.unreduced, a in m.reduced) for a in axes}
    return len(sigs) == 1

When it happens

Trigger: Calling `jax.lax.pcast(x, ('i','j'), to='unreduced')` where x is varying over 'i' but reduced over 'j' (or invarying over one of them).

Common situations: Batching several named axes into one pcast call after mixed collectives; heterogeneous pipelines where different axes were reduced/kept varying by different ops.

Related errors

• Unsupported pcast from={from_}, {to=}
• vary_unreduced_cast only accepts inputs that are varying. Go
• vary_unreduced_cast is a Varying->Unreduced collective. This
• vary_unreduced_cast input cannot be unreduced across the axi
• reduced_vary_cast only accepts inputs that are reduced. Got

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