{"record":{"id":"2c35e1d46f6873c4","repo":"jax-ml/jax","slug":"vary-unreduced-cast-is-a-varying-unreduced-collec","errorCode":null,"errorMessage":"vary_unreduced_cast is a Varying->Unreduced collective. This means that the axis names mentioned in `axes` passed to `vary_unreduced_cast` must be present in `jax.typeof(x).mat.varying`. Got axes={axes} and jax.typeof(x).mat.varying={aval.mat.varying}","messagePattern":"vary_unreduced_cast is a Varying->Unreduced collective\\. This means that the axis names mentioned in `axes` passed to `vary_unreduced_cast` must be present in `jax\\.typeof\\(x\\)\\.mat\\.varying`\\. Got axes=(.+?) and jax\\.typeof\\(x\\)\\.mat\\.varying=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2933,"sourceCode":"  assert set(new_axes) == set(axes)\n  del axes\n  return tree_util.tree_map(\n      lambda leaf: vary_unreduced_cast_p.bind(leaf, axes=new_axes), x)\n\nvary_unreduced_cast_p = core.Primitive('vary_unreduced_cast_p')\nvary_unreduced_cast_p.def_impl(partial(_raise_valueerror, 'vary_unreduced_cast'))\nmlir.register_lowering(vary_unreduced_cast_p, lambda ctx, x, *, axes: [x])\n\ndef _vary_unreduced_cast_abstract_eval(aval, *, axes):\n  assert isinstance(axes, tuple)\n  _check_axis_names(axes, 'vary_unreduced_cast')\n  check_unreduced_args([aval], axes, 'vary_unreduced_cast')\n  if not aval.mat.varying:\n    raise ValueError('vary_unreduced_cast only accepts inputs that are'\n                     f' varying. Got {aval.str_short(True)}')\n  # If the intersection between aval.mat.varying and axes is empty, error\n  if not (aval.mat.varying & set(axes)):\n    raise ValueError(\n        \"vary_unreduced_cast is a Varying->Unreduced collective. This\"\n        \" means that the axis names mentioned in `axes` passed to\"\n        \" `vary_unreduced_cast` must be present in\"\n        f\" `jax.typeof(x).mat.varying`. Got axes={axes} and\"\n        f\" jax.typeof(x).mat.varying={aval.mat.varying}\")\n  if aval.mat.unreduced & set(axes):\n    raise ValueError(\n        \"vary_unreduced_cast input cannot be unreduced across the axis_name\"\n        f\" provided. Got x={aval.str_short(True)} and axis_name={axes}\")\n\n  new_unreduced = aval.mat.unreduced | frozenset(axes)\n  out_vma = frozenset(i for i in aval.mat.varying if i not in axes)\n  return aval.update(manual_axis_type=aval.mat.update(\n    varying=out_vma, unreduced=new_unreduced))\nvary_unreduced_cast_p.def_abstract_eval(_vary_unreduced_cast_abstract_eval)\n\ndef _vary_unreduced_cast_transpose_rule(cts, x, *, axes):\n  assert ad.is_undefined_primal(x)","sourceCodeStart":2915,"sourceCodeEnd":2951,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2915-L2951","documentation":"For `vary_unreduced_cast(x, axes)` the axes you name must actually appear in the input's varying set. The abstract eval intersects `aval.mat.varying` with `set(axes)` and errors when the intersection is empty, telling you exactly which state the input had.","triggerScenarios":"Calling vary_unreduced_cast with axis_name 'i' when x is varying only over 'j' (or over a differently-named axis); typos or renamed axis names after refactoring.","commonSituations":"Renaming mesh/axis names in pjit sharding constraints or Mesh definitions without updating cast call sites; copying cast code between models with different axis naming conventions.","solutions":["Print `jax.typeof(x).mat.varying` and pass one of those exact axis names","Align the axis names with your `jax.sharding.Mesh` / NamedSharding axis names","Replace with `pcast(x, correct_axis, to='unreduced')`"],"exampleFix":"// before\njax.lax.vary_unreduced_cast(x, 'replica')  # x varies over 'data'\n// after\njax.lax.vary_unreduced_cast(x, 'data')","handlingStrategy":"type-guard","validationCode":"t = jax.typeof(x)\nassert set(axes) & t.mat.varying, f'axes {axes} not in varying {t.mat.varying}'","typeGuard":"def axes_in_varying(x, axes) -> bool:\n    return bool(set(axes) & jax.typeof(x).mat.varying)","tryCatchPattern":null,"preventionTips":["Define axis names once as module constants","Mirror Mesh axis_names exactly in cast calls"],"tags":["jax","axis-name","named-axes"],"backgroundTag":"jax-axis-name-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}