jax-ml/jax · error · NotImplementedError

psend is currently only implemented on GPU

Error message

psend is currently only implemented on GPU

What it means

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.

Source

Thrown at jax/_src/lax/parallel.py:1305

effects_lib.lowerable_effects.add_type(SingleSideCollectiveEffect)


def _psend_abstract_eval(x, *, axis_name, **params):
  _check_axis_names(axis_name, 'psend')
  return abstract_token, {
      *map(core.NamedAxisEffect, axis_name),
      single_side_collective_effect,
  }


psend_p = core.Primitive("psend")
psend_p.def_impl(partial(dispatch.apply_primitive, psend_p))
psend_p.def_effectful_abstract_eval(_psend_abstract_eval)
mlir.register_lowering(psend_p, _psend_lowering_gpu, platform="gpu")

def _psend_lowering(ctx, x, *, axis_name, perm):
  raise NotImplementedError("psend is currently only implemented on GPU")
mlir.register_lowering(psend_p, _psend_lowering)

batching.fancy_primitive_batchers[psend_p] = _ppermute_batcher


def _precv_lowering_gpu(ctx, token, *, out_shape, axis_name, perm):
  full_perm, other_args = _pcollectives_lowering_common(
      ctx, axis_name=axis_name, perm=perm, op_name="precv"
  )
  out_type = mlir.aval_to_ir_type(ctx.module_context, out_shape)
  recv_op = hlo.RecvOp(
      [out_type, token.type],
      token,
      source_target_pairs=mlir.dense_int_elements(full_perm),
      **other_args,
  )
  axis_ctx = ctx.module_context.axis_context
  if not isinstance(axis_ctx, SPMDAxisContext):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Target a CUDA/ROCm backend
  2. Add platform dispatch around psend usage
  3. Track upstream CPU/TPU support

Example fix

// before
out = jax.jit(f_with_psend)(x)
// after
assert jax.default_backend() in ('gpu',), 'psend requires GPU'
out = jax.jit(f_with_psend)(x)
Defensive patterns

Strategy: try-catch

Validate before calling

if jax.default_backend() not in ('gpu', 'cuda', 'rocm'):
    raise SystemExit('this script requires a GPU backend for psend')

Type guard

def supports_psend(): return jax.default_backend() in ('gpu','cuda','rocm')

Try / catch

try:
    compile_and_run()
except NotImplementedError as e:
    if 'psend' in str(e): raise RuntimeError('Run on a CUDA/ROCm backend') from e
    raise

Prevention

When it happens

Trigger: 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).

Common situations: Same as 1106: CPU dev machines, TPU pods, GPU-less CI runners.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/4792e2b86c1e3426. Report an issue: GitHub.