{"record":{"id":"29c375df1b033a1b","repo":"jax-ml/jax","slug":"name-can-only-accept-axis-name-which-corresponds","errorCode":null,"errorMessage":"{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}","messagePattern":"(.+?) can only accept axis_name which corresponds to one of varying, unreduced, reduced or invarying state of the input\\. Got input type: (.+?), axes: (.+?) and input state: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":3018,"sourceCode":"  raise NotImplementedError\nbatching.primitive_batchers[core.reduced_vary_cast_p] = _reduced_vary_cast_batcher\n\n################################## pcast #############################\n\ndef _get_from(aval, axes: tuple[AxisName, ...], name) -> str:\n  out = set()\n  for a in axes:\n    if a in aval.mat.varying:\n      out.add('varying')\n    elif a in aval.mat.unreduced:\n      out.add('unreduced')\n    elif a in aval.mat.reduced:\n      out.add('reduced')\n    else:\n      out.add('invarying')\n\n  if len(out) > 1:\n    raise ValueError(\n        f\"{name} can only accept axis_name which corresponds to one of\"\n        \" varying, unreduced, reduced or invarying state of the input. Got\"\n        f\" input type: {aval}, axes: {axes} and input state: {out}\")\n  o, = out\n  return o\n\n\n_pcast_funcs = {\n    ('invarying', 'varying'): core.pvary,\n    ('invarying', 'reduced'): preduced,\n    ('varying', 'unreduced'): vary_unreduced_cast,\n    ('reduced', 'varying'): core.reduced_vary_cast,\n}\n\n_allowed_pcast_to = {'unreduced', 'reduced', 'varying'}\n\ndef pcast(x, axis_name, *, to: str):\n  if isinstance(axis_name, (set, frozenset)):","sourceCodeStart":3000,"sourceCodeEnd":3036,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L3000-L3036","documentation":"`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.","triggerScenarios":"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).","commonSituations":"Batching several named axes into one pcast call after mixed collectives; heterogeneous pipelines where different axes were reduced/kept varying by different ops.","solutions":["Split the pcast into one call per axis (or per group of axes sharing the same state)","First normalize the input to a single state using the explicit casts (vary_unreduced_cast/reduced_vary_cast/...) per axis","Inspect `jax.typeof(x).mat` and re-order preceding collectives so all requested axes share one state"],"exampleFix":"# before\ny = jax.lax.pcast(x, ('i', 'j'), to='unreduced')\n# after\ny = jax.lax.pcast(jax.lax.pcast(x, 'i', to='unreduced'), 'j', to='unreduced')","handlingStrategy":"validation","validationCode":"states = {jax.typeof(x).mat.state_of(a) for a in axes}  # conceptual\n# practically: check each axis's membership in varying/unreduced/reduced sets and group","typeGuard":"def axes_share_state(x, axes) -> bool:\n    m = jax.typeof(x).mat\n    sigs = {(a in m.varying, a in m.unreduced, a in m.reduced) for a in axes}\n    return len(sigs) == 1","tryCatchPattern":null,"preventionTips":["pcast one axis at a time when states may differ"],"tags":["jax","pcast","named-axes","mixed-state"],"backgroundTag":"jax-manual-axis-state-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}