{"record":{"id":"522e9506d3fa9a0a","repo":"jax-ml/jax","slug":"precv-is-currently-only-implemented-on-gpu","errorCode":null,"errorMessage":"precv is currently only implemented on GPU","messagePattern":"precv is currently only implemented on GPU","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1343,"sourceCode":"\n  # recv_op should return an array of [RankedTensorType, StableHlo.token]; we\n  # only need the tensor.\n  results = recv_op.results\n  return [results[0]]\n\n\ndef _precv_abstract_eval(\n    token, *, out_shape, axis_name, **params\n):\n  return out_shape, {*map(core.NamedAxisEffect, axis_name),\n                     single_side_collective_effect}\n\nprecv_p = core.Primitive(\"precv\")\nprecv_p.def_effectful_abstract_eval(_precv_abstract_eval)\nmlir.register_lowering(precv_p, _precv_lowering_gpu, platform='gpu')\n\ndef _precv_lowering(ctx, token, *, out_shape, axis_name, perm):\n  raise NotImplementedError(\"precv is currently only implemented on GPU\")\nmlir.register_lowering(precv_p, _precv_lowering)\n\nbatching.fancy_primitive_batchers[precv_p] = _ppermute_batcher\n\ndef _pbroadcast_transpose_rule(t, x, source, axis_name):\n  is_source = axis_index(axis_name) == source\n  tsum = psum(t, axis_name)\n  return [lax.select(is_source, lax.full_like(t, tsum), lax.full_like(t, 0))]\n\ndef _pbroadcast_batcher(axis_data, vals_in, dims_in, axis_name, source):\n  axis_size = axis_data.size\n  (v,), (d,) = vals_in, dims_in\n  if not isinstance(axis_name, (tuple, list)):\n    axis_name = (axis_name,)\n  if d is None and axis_data.name not in axis_name:\n    return pbroadcast_p.bind(v, axis_name=axis_name, source=source), None\n  if axis_data.name not in axis_name:\n    return pbroadcast_p.bind(v, axis_name=axis_name, source=source), d","sourceCodeStart":1325,"sourceCodeEnd":1361,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1325-L1361","documentation":"Fallback lowering for precv_p on non-GPU platforms: precv only has a GPU (CUDA/ROCm) lowering rule, so compiling it elsewhere hits this unconditional NotImplementedError with a clear message.","triggerScenarios":"Compiling code containing lax.precv on CPU or TPU.","commonSituations":"CPU local testing of distributed pipelines; GPU-less CI.","solutions":["Run on CUDA/ROCm backends","Gate precv paths on backend capability","Provide a non-NCCL fallback (e.g. all_gather based) for CPU tests"],"exampleFix":"// before\ny = jax.jit(f_with_precv)(token)\n// after\nif jax.default_backend() != 'gpu':\n    y = mock_recv_fallback(token)\nelse:\n    y = jax.jit(f_with_precv)(token)","handlingStrategy":"type-guard","validationCode":"if jax.default_backend() not in ('gpu', 'cuda', 'rocm'):\n    pytest.skip('precv requires GPU')","typeGuard":"def supports_precv(): return jax.default_backend() in ('gpu','cuda','rocm')","tryCatchPattern":"try:\n    jax.jit(f)(x)\nexcept NotImplementedError as e:\n    if 'precv' in str(e): y = mock_pipeline(x)\n    else: raise","preventionTips":["Gate distributed point-to-point tests on GPU availability"],"tags":["jax","precv","backend","not-implemented"],"backgroundTag":"unsupported-backend","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}