{"record":{"id":"2a544c573eb8fa4f","repo":"jax-ml/jax","slug":"vary-unreduced-cast-input-cannot-be-unreduced-acro","errorCode":null,"errorMessage":"vary_unreduced_cast input cannot be unreduced across the axis_name provided. Got x={aval.str_short(True)} and axis_name={axes}","messagePattern":"vary_unreduced_cast input cannot be unreduced across the axis_name provided\\. Got x=(.+?) and axis_name=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2940,"sourceCode":"mlir.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)\n  return (core.reduced_vary_cast(cts, axis_name=axes),)\nad.deflinear2(vary_unreduced_cast_p, _vary_unreduced_cast_transpose_rule)\n\ndef _vary_unreduced_cast_batcher(vals_in, dims_in, *, axes):\n  raise NotImplementedError\nbatching.primitive_batchers[vary_unreduced_cast_p] = _vary_unreduced_cast_batcher\n","sourceCodeStart":2922,"sourceCodeEnd":2958,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2922-L2958","documentation":"`vary_unreduced_cast` cannot mark an axis unreduced if the input is already unreduced across that axis. The abstract eval rejects inputs where `aval.mat.unreduced & set(axes)` is non-empty, since the cast would be redundant/contradictory.","triggerScenarios":"Applying vary_unreduced_cast to the output of a previous vary_unreduced_cast (or unreduced_psum) over the same axis_name; double-casting in composed collective helpers.","commonSituations":"Wrapper functions that defensively cast to unreduced being applied to already-cast values; layered abstractions over manual collectives stacking casts.","solutions":["Remove the duplicate cast — the value is already unreduced across that axis","Guard the cast with a check on `jax.typeof(x).mat.unreduced`","Use `pcast`, which is idempotent-aware dispatch, instead of the raw cast"],"exampleFix":"// before\nx = vary_unreduced_cast(x, 'dev')\nx = vary_unreduced_cast(x, 'dev')  # duplicate\n// after\nx = vary_unreduced_cast(x, 'dev')","handlingStrategy":"type-guard","validationCode":"if set(axes) & jax.typeof(x).mat.unreduced:\n    return x  # already unreduced; skip cast","typeGuard":"def already_unreduced(x, axes) -> bool:\n    return bool(set(axes) & jax.typeof(x).mat.unreduced)","tryCatchPattern":null,"preventionTips":["Make cast wrappers idempotent by checking mat state first","Prefer pcast over raw casts"],"tags":["jax","named-axes","duplicate-cast"],"backgroundTag":"jax-manual-axis-state-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}