{"record":{"id":"60961ca5063d1383","repo":"jax-ml/jax","slug":"psend-is-currently-only-implemented-on-gpus","errorCode":null,"errorMessage":"psend is currently only implemented on GPUs","messagePattern":"psend is currently only implemented on GPUs","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1269,"sourceCode":"batching.fancy_primitive_batchers[ppermute_p] = _ppermute_batcher\n\n\n@dataclass(frozen=True, slots=True)\nclass SingleSideCollectiveEffect(core.Effect):\n  __str__ = lambda _: \"one-sided communication\"\n  def __hash__(self):\n    return hash(SingleSideCollectiveEffect)\n  def __eq__(self, other):\n    return isinstance(other, SingleSideCollectiveEffect)\n\n\nsingle_side_collective_effect = SingleSideCollectiveEffect()\ncore.effects.control_flow_allowed_effects.add_type(SingleSideCollectiveEffect)\n\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","sourceCodeStart":1251,"sourceCodeEnd":1287,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1251-L1287","documentation":"psend is a point-to-point collective whose lowering only exists for CUDA/ROCm platforms. On any other backend (CPU, TPU) the GPU lowering rule raises NotImplementedError.","triggerScenarios":"Compiling a function containing jax.lax.psend on CPU or TPU (jax.default_device or available backends).","commonSituations":"Developing/testing distributed code locally on CPU before running on GPU clusters; CI without GPUs hitting the psend path.","solutions":["Run on a CUDA or ROCm backend (set jax.default_device to a GPU or run on GPU hosts)","Guard psend/precv usage with jax.devices() platform checks and provide a CPU fallback","Request/await CPU support upstream"],"exampleFix":"// before\ny = jax.jit(lambda x, t: lax.psend(x, t, 'i', perm))(x, token)\n// after\nif jax.default_backend() in ('gpu', 'cuda', 'rocm'):\n    y = jax.jit(lambda x, t: lax.psend(x, t, 'i', perm))(x, token)\nelse:\n    y = cpu_fallback(x)","handlingStrategy":"type-guard","validationCode":"import jax\nassert jax.default_backend() in ('gpu', 'cuda', 'rocm'), 'psend requires a GPU backend'","typeGuard":"def has_gpu_backend():\n    return jax.default_backend() in ('gpu', 'cuda', 'rocm')","tryCatchPattern":"try:\n    y = jax.jit(f)(x)\nexcept NotImplementedError as e:\n    if 'psend' in str(e): y = cpu_stub(x)\n    else: raise","preventionTips":["Skip GPU-only collective tests when jax.default_backend() == 'cpu'","Mark psend tests with pytest.mark.gpu"],"tags":["jax","psend","gpu-only","backend"],"backgroundTag":"unsupported-backend","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}