{"record":{"id":"713e023cc0e77d10","repo":"jax-ml/jax","slug":"precv-currently-only-supports-manual-sharding","errorCode":null,"errorMessage":"precv currently only supports manual sharding","messagePattern":"precv currently only supports manual sharding","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1324,"sourceCode":"mlir.register_lowering(psend_p, _psend_lowering)\n\nbatching.fancy_primitive_batchers[psend_p] = _ppermute_batcher\n\n\ndef _precv_lowering_gpu(ctx, token, *, out_shape, axis_name, perm):\n  full_perm, other_args = _pcollectives_lowering_common(\n      ctx, axis_name=axis_name, perm=perm, op_name=\"precv\"\n  )\n  out_type = mlir.aval_to_ir_type(ctx.module_context, out_shape)\n  recv_op = hlo.RecvOp(\n      [out_type, token.type],\n      token,\n      source_target_pairs=mlir.dense_int_elements(full_perm),\n      **other_args,\n  )\n  axis_ctx = ctx.module_context.axis_context\n  if not isinstance(axis_ctx, SPMDAxisContext):\n    raise NotImplementedError(\"precv currently only supports manual sharding\")\n\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):","sourceCodeStart":1306,"sourceCodeEnd":1342,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1306-L1342","documentation":"precv's GPU lowering requires an SPMDAxisContext: the call must be under jax.shard_map with manual sharding so that receiver replica semantics are defined. Other axis contexts (automatic sharding, eager) are rejected.","triggerScenarios":"Calling jax.lax.precv outside jax.shard_map, or under jit with automatic (GSD) sharding.","commonSituations":"Building pipeline-parallel send/recv pairs and forgetting the shard_map wrapper on the receiving side.","solutions":["Wrap precv (and matching psend) in jax.shard_map over a mesh","Give inputs a NamedSharding matching the mesh","Verify with jax.debug that the axis context is SPMD before calling"],"exampleFix":"// before\ny = lax.precv(token, shape, 'i', perm)\n// after\ny = jax.shard_map(lambda tok: lax.precv(tok, shape, 'i', perm), mesh)(token)","handlingStrategy":"validation","validationCode":"null  # structural: ensure precv appears only inside shard_map bodies","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Wrap send/recv pairs in jax.shard_map","Test on small meshes with manual sharding first"],"tags":["jax","precv","shard-map","manual-sharding"],"backgroundTag":"missing-context","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}