{"record":{"id":"ea6f55d6d19cec16","repo":"jax-ml/jax","slug":"reduced-vary-cast-only-accepts-inputs-that-are-red","errorCode":null,"errorMessage":"reduced_vary_cast only accepts inputs that are reduced. Got {aval.str_short(True)}","messagePattern":"reduced_vary_cast only accepts inputs that are reduced\\. Got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2971,"sourceCode":"ad.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\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)","sourceCodeStart":2953,"sourceCodeEnd":2989,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2953-L2989","documentation":"`reduced_vary_cast` is a Reduced->Varying no-op cast; its abstract eval requires the input to currently be in the reduced state (`aval.mat.reduced` non-empty). Passing a value that is varying/unreduced/invarying raises this ValueError.","triggerScenarios":"Calling jax.lax.core reduced_vary_cast (exposed via pcast dispatch) on a value produced without a reducing collective (e.g. a plain array or an unreduced output).","commonSituations":"Hand-managing cast chains between axis states and calling the wrong directional cast; assuming a psum was applied when it was optimized away.","solutions":["Inspect `jax.typeof(x).mat.reduced` and only cast values that are reduced","Use `jax.lax.pcast(x, axis_name, to='varying')` which picks the right cast for the actual state","Apply the intended reduction (psum etc.) before the cast"],"exampleFix":"// before\ny = reduced_vary_cast(x, 'dev')  # x not reduced\n// after\ny = pcast(psum_like(x, 'dev'), 'dev', to='varying')\n# or simply pcast(x, 'dev', to='varying')","handlingStrategy":"type-guard","validationCode":"if not jax.typeof(x).mat.reduced:\n    x = apply_reduction(x, axes)  # real collective first","typeGuard":"def is_reduced(x) -> bool:\n    return bool(jax.typeof(x).mat.reduced)","tryCatchPattern":"catch ValueError, fall back to pcast(x, axes, to='varying')","preventionTips":["Use pcast for state-aware dispatch","Assert the reduced state right after each reducing collective"],"tags":["jax","named-axes","manual-parallelism","aval-validation"],"backgroundTag":"jax-manual-axis-state-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}