{"record":{"id":"4792e2b86c1e3426","repo":"jax-ml/jax","slug":"psend-is-currently-only-implemented-on-gpu","errorCode":null,"errorMessage":"psend is currently only implemented on GPU","messagePattern":"psend is currently only implemented on GPU","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1305,"sourceCode":"\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)\nmlir.register_lowering(psend_p, _psend_lowering_gpu, platform=\"gpu\")\n\ndef _psend_lowering(ctx, x, *, axis_name, perm):\n  raise NotImplementedError(\"psend is currently only implemented on GPU\")\nmlir.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):","sourceCodeStart":1287,"sourceCodeEnd":1323,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1287-L1323","documentation":"This is the fallback (non-GPU) lowering rule registered for psend_p; it exists so compilation fails with a clear message instead of an opaque 'no lowering' error whenever psend is compiled on a non-GPU platform.","triggerScenarios":"Any compilation of a trace containing lax.psend on CPU or TPU backends (the gpu-platform rule doesn't match, so the default rule fires).","commonSituations":"Same as 1106: CPU dev machines, TPU pods, GPU-less CI runners.","solutions":["Target a CUDA/ROCm backend","Add platform dispatch around psend usage","Track upstream CPU/TPU support"],"exampleFix":"// before\nout = jax.jit(f_with_psend)(x)\n// after\nassert jax.default_backend() in ('gpu',), 'psend requires GPU'\nout = jax.jit(f_with_psend)(x)","handlingStrategy":"try-catch","validationCode":"if jax.default_backend() not in ('gpu', 'cuda', 'rocm'):\n    raise SystemExit('this script requires a GPU backend for psend')","typeGuard":"def supports_psend(): return jax.default_backend() in ('gpu','cuda','rocm')","tryCatchPattern":"try:\n    compile_and_run()\nexcept NotImplementedError as e:\n    if 'psend' in str(e): raise RuntimeError('Run on a CUDA/ROCm backend') from e\n    raise","preventionTips":["Check backend before building graphs with psend"],"tags":["jax","psend","backend","not-implemented"],"backgroundTag":"unsupported-backend","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}