{"record":{"id":"e06f3fc14a217fb7","repo":"jax-ml/jax","slug":"vary-unreduced-cast-only-accepts-inputs-that-are-v","errorCode":null,"errorMessage":"vary_unreduced_cast only accepts inputs that are varying. Got {aval.str_short(True)}","messagePattern":"vary_unreduced_cast only accepts inputs that are varying\\. Got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2929,"sourceCode":"  cur_mesh = get_abstract_mesh()\n  if not config._check_vma.value and all(a in cur_mesh.manual_axes for a in axes):\n    return x\n  new_axes = axes if cur_mesh.empty else core.order_wrt_mesh(cur_mesh, axes)\n  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))","sourceCodeStart":2911,"sourceCodeEnd":2947,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2911-L2947","documentation":"`vary_unreduced_cast` is a Varying->Unreduced no-op cast used in JAX's manual-parallelism (named axis) machinery. Its abstract eval requires the input aval to be marked as varying (`aval.mat.varying` non-empty). Passing a value that is not varying raises this ValueError.","triggerScenarios":"Calling `jax.lax.vary_unreduced_cast(x, axis_name)` where `jax.typeof(x).mat.varying` is empty — e.g. the input was already fully reduced/unreduced/invarying across all named axes.","commonSituations":"Hand-writing cast sequences between manual axis states (varying/unreduced/reduced) and losing track of the current state; mixing pmap-style code with collectives that implicitly reduce, then casting the result.","solutions":["Check `jax.typeof(x).mat` before casting and only call vary_unreduced_cast on values whose `varying` set is non-empty","Use `jax.lax.pcast(x, axis_name, to='unreduced')`, which dispatches to the correct cast for the actual input state","Restructure the collective sequence so the operand is genuinely varying (e.g. re-materialize variation with a varying cast/split before)"],"exampleFix":"// before\ny = jax.lax.vary_unreduced_cast(x, 'dev')  # x not varying\n// after\ny = jax.lax.pcast(x, 'dev', to='unreduced')","handlingStrategy":"type-guard","validationCode":"t = jax.typeof(x)\nif not t.mat.varying:\n    x = to_varying_state(x, axes)  # your normalization helper","typeGuard":"def is_varying(x, ) -> bool:\n    return bool(jax.typeof(x).mat.varying)","tryCatchPattern":"catch ValueError and re-dispatch via pcast(x, axes, to='unreduced')","preventionTips":["Route all casts through jax.lax.pcast rather than the raw primitives","Log jax.typeof(x).mat at collective boundaries during development"],"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"}