{"record":{"id":"519f29f51aca4e7b","repo":"jax-ml/jax","slug":"psend-currently-only-supports-manual-sharding","errorCode":null,"errorMessage":"psend currently only supports manual sharding","messagePattern":"psend currently only supports manual sharding","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1283,"sourceCode":"\ndef _psend_lowering_gpu(ctx, x, *, axis_name, perm):\n  if (\"cuda\" not in ctx.module_context.platforms and\n      \"rocm\" not in ctx.module_context.platforms):\n    raise NotImplementedError(\"psend is currently only implemented on GPUs\")\n\n  full_perm, other_args = _pcollectives_lowering_common(\n      ctx, axis_name=axis_name, perm=perm, op_name=\"psend\"\n  )\n  token = hlo.create_token()\n  send_op = hlo.SendOp(\n      [x],\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(\"psend currently only supports manual sharding\")\n\n  return send_op.results\n\n\neffects_lib.lowerable_effects.add_type(SingleSideCollectiveEffect)\n\n\ndef _psend_abstract_eval(x, *, axis_name, **params):\n  _check_axis_names(axis_name, 'psend')\n  return abstract_token, {\n      *map(core.NamedAxisEffect, axis_name),\n      single_side_collective_effect,\n  }\n\n\npsend_p = core.Primitive(\"psend\")\npsend_p.def_impl(partial(dispatch.apply_primitive, psend_p))\npsend_p.def_effectful_abstract_eval(_psend_abstract_eval)","sourceCodeStart":1265,"sourceCodeEnd":1301,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1265-L1301","documentation":"psend's lowering requires an SPMDAxisContext, i.e. the code must be under jax.shard_map with manual sharding so replica semantics apply. Under automatic (GSD) sharding or eager execution the axis context is a different type and lowering is refused.","triggerScenarios":"Calling jax.lax.psend outside jax.shard_map (e.g. jit with automatic sharding, or traced eagerly).","commonSituations":"Writing point-to-point communication pipelines and forgetting the shard_map wrapper; assuming jit+NamedSharding provides manual sharding.","solutions":["Wrap the psend/precv code in jax.shard_map with mesh and manual sharding","Ensure the sharded inputs have a sharding compatible with the mesh (e.g. NamedSharding(mesh, ...))","Move communication logic into a shard_map function passed to jit"],"exampleFix":"// before\ny = jax.jit(lambda x, t: lax.psend(x, t, 'i', perm))(x, token)\n// after\ny = jax.jit(jax.shard_map(lambda x, t: lax.psend(x, t, 'i', perm), mesh))(x, token)","handlingStrategy":"validation","validationCode":"assert in_shard_map, 'psend must be called inside jax.shard_map (manual sharding)'  # enforce via code review/lint","typeGuard":null,"tryCatchPattern":"null","preventionTips":["Keep all psend/precv calls inside shard_map-wrapped functions","Pair every psend with a matching precv in the same shard_map scope"],"tags":["jax","psend","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"}