jax-ml/jax · error · ValueError

Unsupported pcast from={from_}, {to=}

Error message

Unsupported pcast from={from_}, {to=}

What it means

`pcast` looks up a cast function in `_pcast_funcs` keyed by (from_state, to). If no cast primitive is implemented for that transition (e.g. from 'invarying' or an 'unreduced'->'reduced' pair not in the table), it raises this ValueError.

Source

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

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

def pcast(x, axis_name, *, to: str):
  if isinstance(axis_name, (set, frozenset)):
    raise TypeError(f"{axis_name=} must be a tuple or a str. Got {axis_name}")
  axes = (axis_name,) if not isinstance(axis_name, tuple) else axis_name
  if not axis_name:
    return x

  if to not in _allowed_pcast_to:
    raise ValueError(
        "Got unexpected `to` value. Allowed `to` values are:"
        f" {_allowed_pcast_to}")

  def bind(leaf):
    from_ = _get_from(core.typeof(leaf), axes, 'jax.lax.pcast')
    func = _pcast_funcs.get((from_, to), None)
    if func is None:
      raise ValueError(f"Unsupported pcast from={from_}, {to=}")
    return func(leaf, axes)
  return tree_util.tree_map(bind, x)

def _emit_async_start_custom_call(
    target_name, ctx, x, cfg, called_computations=None
):
  out_aval, = ctx.avals_out
  future_type = mlir.aval_to_ir_type(ctx.module_context, out_aval.inner_aval)

  cfg = dict(cfg)
  if "channel_handle" in cfg:
    cfg["channel_id"] = cfg.pop("channel_handle").handle
  if "use_global_device_ids" in cfg:
    cfg["use_global_device_ids"] = cfg["use_global_device_ids"].value

  def _json_default(obj):
    if isinstance(obj, np.integer):
      return int(obj)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect `jax.typeof(x).mat` and insert the appropriate collective first (e.g. an actual psum for unreduced->reduced semantics)
2. Choose a supported transition or chain two supported casts
3. Check the `_pcast_funcs` table keys in your JAX version to see which transitions exist

Example fix

# before
y = pcast(x, 'dev', to='reduced')  # x invarying over 'dev'
# after
y = psum_like_reduce(x, 'dev')  # apply real reduction, then cast as needed

m = jax.typeof(x).mat
# ensure the (from,to) pair is supported before calling; from must be varying/unreduced/reduced with an entry in _pcast_funcs

def pcast_supported(x, axes, to) -> bool:
    m = jax.typeof(x).mat
    for a in axes:
        if not (a in m.varying or a in m.unreduced or a in m.reduced):
            return False
    return to in {'unreduced', 'reduced', 'varying'}

catch ValueError and apply an explicit collective instead of pcast

When it happens

Trigger: Calling `pcast(x, 'dev', to='reduced')` when x is invarying or unreduced over 'dev' — transitions without a defined collective.

Common situations: Trying to use pcast as a general state converter; assuming every (from,to) pair exists because the target `to` was validated.

Related errors

• {name} can only accept axis_name which corresponds to one of
• 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/e9bb3d3424420edd. Report an issue: GitHub.