{"record":{"id":"e9bb3d3424420edd","repo":"jax-ml/jax","slug":"unsupported-pcast-from-from-to","errorCode":null,"errorMessage":"Unsupported pcast from={from_}, {to=}","messagePattern":"Unsupported pcast from=(.+?), (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":3051,"sourceCode":"_allowed_pcast_to = {'unreduced', 'reduced', 'varying'}\n\ndef pcast(x, axis_name, *, to: str):\n  if isinstance(axis_name, (set, frozenset)):\n    raise TypeError(f\"{axis_name=} must be a tuple or a str. Got {axis_name}\")\n  axes = (axis_name,) if not isinstance(axis_name, tuple) else axis_name\n  if not axis_name:\n    return x\n\n  if to not in _allowed_pcast_to:\n    raise ValueError(\n        \"Got unexpected `to` value. Allowed `to` values are:\"\n        f\" {_allowed_pcast_to}\")\n\n  def bind(leaf):\n    from_ = _get_from(core.typeof(leaf), axes, 'jax.lax.pcast')\n    func = _pcast_funcs.get((from_, to), None)\n    if func is None:\n      raise ValueError(f\"Unsupported pcast from={from_}, {to=}\")\n    return func(leaf, axes)\n  return tree_util.tree_map(bind, x)\n\ndef _emit_async_start_custom_call(\n    target_name, ctx, x, cfg, called_computations=None\n):\n  out_aval, = ctx.avals_out\n  future_type = mlir.aval_to_ir_type(ctx.module_context, out_aval.inner_aval)\n\n  cfg = dict(cfg)\n  if \"channel_handle\" in cfg:\n    cfg[\"channel_id\"] = cfg.pop(\"channel_handle\").handle\n  if \"use_global_device_ids\" in cfg:\n    cfg[\"use_global_device_ids\"] = cfg[\"use_global_device_ids\"].value\n\n  def _json_default(obj):\n    if isinstance(obj, np.integer):\n      return int(obj)","sourceCodeStart":3033,"sourceCodeEnd":3069,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L3033-L3069","documentation":"`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.","triggerScenarios":"Calling `pcast(x, 'dev', to='reduced')` when x is invarying or unreduced over 'dev' — transitions without a defined collective.","commonSituations":"Trying to use pcast as a general state converter; assuming every (from,to) pair exists because the target `to` was validated.","solutions":["Inspect `jax.typeof(x).mat` and insert the appropriate collective first (e.g. an actual psum for unreduced->reduced semantics)","Choose a supported transition or chain two supported casts","Check the `_pcast_funcs` table keys in your JAX version to see which transitions exist"],"exampleFix":"# before\ny = pcast(x, 'dev', to='reduced')  # x invarying over 'dev'\n# after\ny = psum_like_reduce(x, 'dev')  # apply real reduction, then cast as needed","handlingStrategy":"validation","validationCode":"m = jax.typeof(x).mat\n# ensure the (from,to) pair is supported before calling; from must be varying/unreduced/reduced with an entry in _pcast_funcs","typeGuard":"def pcast_supported(x, axes, to) -> bool:\n    m = jax.typeof(x).mat\n    for a in axes:\n        if not (a in m.varying or a in m.unreduced or a in m.reduced):\n            return False\n    return to in {'unreduced', 'reduced', 'varying'}","tryCatchPattern":"catch ValueError and apply an explicit collective instead of pcast","preventionTips":["Remember pcast is a cast, not a reducer; perform real reductions with psum first"],"tags":["jax","pcast","unsupported-transition","named-axes"],"backgroundTag":"unsupported-operation-transition","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}