{"record":{"id":"fc44b754b81d81e8","repo":"jax-ml/jax","slug":"pvary-is-a-invariant-variant-collective-this-mea","errorCode":null,"errorMessage":"pvary is a invariant->variant collective. This means that the axis names mentioned in `axes` passed to `pvary` must not be present in `jax.typeof(inp).mat.varying`. Got axes={axes} and jax.typeof(inp)={aval}","messagePattern":"pvary is a invariant->variant collective\\. This means that the axis names mentioned in `axes` passed to `pvary` must not be present in `jax\\.typeof\\(inp\\)\\.mat\\.varying`\\. Got axes=(.+?) and jax\\.typeof\\(inp\\)=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":2538,"sourceCode":"def _psum_invariant_transpose_rule(cts, arg, *, axes):\n  assert ad.is_undefined_primal(arg)\n  return (core.pvary(cts, axis_name=axes),)\nad.deflinear2(psum_invariant_p, _psum_invariant_transpose_rule)\n\n########################### pvary ##################################\n\ndef _raise_valueerror(name, arg, *, axes):\n  raise ValueError(f'{name} should be called under jax.shard_map.')\n\ncore.pvary_p.def_impl(partial(_raise_valueerror, 'pvary'))\nmlir.register_lowering(core.pvary_p, lambda ctx, x, *, axes: [x])\n\ndef _pvary_abstract_eval(aval, *, axes):\n  _check_axis_names(axes, 'pvary')\n  check_unreduced_args([aval], axes, 'pvary')\n  assert isinstance(axes, tuple)\n  if set(axes).intersection(aval.mat.varying):\n    raise ValueError(\n        \"pvary is a invariant->variant collective. This means that the axis\"\n        \" names mentioned in `axes` passed to `pvary` must not be present in\"\n        f\" `jax.typeof(inp).mat.varying`. Got axes={axes} and\"\n        f\" jax.typeof(inp)={aval}\")\n  out_vma = aval.mat.varying.union(frozenset(axes))\n  return aval.update(sharding=aval.sharding.update(mesh=get_abstract_mesh()),\n                     manual_axis_type=aval.mat.update(varying=out_vma))\ncore.pvary_p.def_abstract_eval(_pvary_abstract_eval)\n\ndef _pvary_transpose_rule(cts, arg, *, axes):\n  assert ad.is_undefined_primal(arg)\n  return (psum_invariant_p.bind(cts, axes=axes),)\nad.deflinear2(core.pvary_p, _pvary_transpose_rule)\n\ndef _pvary_batcher(vals_in, dims_in, *, axes):\n  if any(type(axis) is int for axis in axes):\n    raise NotImplementedError\n  (x,), (d,) = vals_in, dims_in","sourceCodeStart":2520,"sourceCodeEnd":2556,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L2520-L2556","documentation":"`pvary` marks an array as varying along named mesh axes; it is an invariant->variant collective, so JAX raises this error when any axis in `axes` is already present in the array's `mat.varying` set. Re-varying an already-varying axis is redundant and usually indicates a misunderstanding of the mesh annotation or a duplicated call. The abstract eval in `_pvary_abstract_eval` rejects it before any computation runs.","triggerScenarios":"Calling `lax.pvary(x, 'i')` twice with the same axis name; or passing an axis name that is already in `jax.typeof(x).mat.varying` because a previous pvary/pvarying operation or a collective produced it.","commonSituations":"Refactoring older `pmap`/axis-index code to the new named-mesh `pvary` API and accidentally applying pvary after another primitive that already varied that axis; stacking collectives in a pipeline where the output is already varying.","solutions":["Remove the redundant `pvary` call for axes already listed in `jax.typeof(inp).mat.varying`","Inspect `jax.typeof(x).mat.varying` before calling pvary and only pass axes not present","Check upstream collectives (e.g. a previous pvary or all_gather) that may already vary the axis","If you intended a different semantic, use the appropriate collective (e.g. pbroadcast for the inverse direction)"],"exampleFix":"// before\ny = lax.pvary(lax.pvary(x, 'i'), 'i')\n// after\ny = lax.pvary(x, 'i')","handlingStrategy":"validation","validationCode":"varying = jax.typeof(x).mat.varying\nnew_axes = tuple(a for a in axes if a not in varying)\ny = lax.pvary(x, new_axes) if new_axes else x","typeGuard":"def can_pvary(x, axes: tuple[str, ...]) -> bool:\n    return not (set(axes) & jax.typeof(x).mat.varying)","tryCatchPattern":"try:\n    y = lax.pvary(x, axes)\nexcept ValueError as e:\n    if 'pvary is a invariant->variant' in str(e):\n        y = x  # already varying on those axes\n    else:\n        raise","preventionTips":["Print jax.typeof(x).mat.varying when composing collectives","Keep pvary calls at data-ingress points only, apply each axis once"],"tags":["jax","collectives","mesh","spmd","pvary"],"backgroundTag":"jax-mat-axis-state-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}