{"record":{"id":"d3f50e5140c060dd","repo":"jax-ml/jax","slug":"reduced-vary-cast-is-a-reduced-varying-collective","errorCode":null,"errorMessage":"reduced_vary_cast is a Reduced->Varying collective. This means that the axis names mentioned in `axes` passed to `reduced_vary_cast` must be present in `jax.typeof(x).mat.reduced`. Got axes={axes} and jax.typeof(x).mat.reduced={aval.mat.reduced}","messagePattern":"reduced_vary_cast is a Reduced->Varying collective\\. This means that the axis names mentioned in `axes` passed to `reduced_vary_cast` must be present in `jax\\.typeof\\(x\\)\\.mat\\.reduced`\\. Got axes=(.+?) and jax\\.typeof\\(x\\)\\.mat\\.reduced=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2975,"sourceCode":"batching.primitive_batchers[vary_unreduced_cast_p] = _vary_unreduced_cast_batcher\n\n####################### reduced_vary_cast #############################\n\n# Reduced -> Varying no-op cast\n# Traceable defined in core.py to avoid circular imports\ncore.reduced_vary_cast_p.def_impl(\n    partial(_raise_valueerror, 'reduced_vary_cast'))\nmlir.register_lowering(core.reduced_vary_cast_p, lambda ctx, x, *, axes: [x])\n\ndef _reduced_vary_cast_abstract_eval(aval, *, axes):\n  assert isinstance(axes, tuple)\n  _check_axis_names(axes, 'reduced_vary_cast')\n  if not aval.mat.reduced:\n    raise ValueError('reduced_vary_cast only accepts inputs that are'\n                     f' reduced. Got {aval.str_short(True)}')\n  # If the intersection between aval.mat.reduced and axes is empty, error\n  if not (aval.mat.reduced & set(axes)):\n    raise ValueError(\n        \"reduced_vary_cast is a Reduced->Varying collective. This\"\n        \" means that the axis names mentioned in `axes` passed to\"\n        \" `reduced_vary_cast` must be present in\"\n        f\" `jax.typeof(x).mat.reduced`. Got axes={axes} and\"\n        f\" jax.typeof(x).mat.reduced={aval.mat.reduced}\")\n  if aval.mat.varying & set(axes):\n    raise ValueError(\n        \"reduced_vary_cast input cannot be varying across the axis_name\"\n        f\" provided. Got x={aval.str_short(True)} and axis_name={axes}\")\n  if aval.mat.unreduced:\n    check_unreduced_kind('reduced_vary_cast', aval.mat, UnreducedKind.sum)\n\n  new_reduced = frozenset(i for i in aval.mat.reduced if i not in axes)\n  out_vma = aval.mat.varying | frozenset(axes)\n  return aval.update(manual_axis_type=aval.mat.update(\n    varying=out_vma, reduced=new_reduced))\ncore.reduced_vary_cast_p.def_abstract_eval(_reduced_vary_cast_abstract_eval)\n","sourceCodeStart":2957,"sourceCodeEnd":2993,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2957-L2993","documentation":"For `reduced_vary_cast(x, axes)` every named axis in `axes` must be present in the input's `mat.reduced` set. The abstract eval raises when `aval.mat.reduced & set(axes)` is empty and reports both sets.","triggerScenarios":"Casting to varying over axis 'i' when x is reduced over 'j'; using an axis name that was never reduced over (e.g. the value was reduced under a different mesh axis name).","commonSituations":"Renaming mesh axes or partition specs without updating cast sites; mixing collectives applied under different axis names in one pipeline.","solutions":["Check `jax.typeof(x).mat.reduced` and pass one of those axis names","Ensure the preceding psum/reduce collective used the same axis_name","Use `pcast` with the correct axis and let it dispatch"],"exampleFix":"// before\nreduced_vary_cast(x, 'data')  # reduced over 'model'\n// after\nreduced_vary_cast(x, 'model')","handlingStrategy":"type-guard","validationCode":"t = jax.typeof(x)\nassert set(axes) & t.mat.reduced, f'axes {axes} not reduced; reduced={t.mat.reduced}'","typeGuard":"def axes_in_reduced(x, axes) -> bool:\n    return bool(set(axes) & jax.typeof(x).mat.reduced)","tryCatchPattern":null,"preventionTips":["Use the same axis_name string in the reduce and the cast","Centralize axis-name constants"],"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"}