jax-ml/jax · error · NotImplementedError

psend is currently only implemented on GPUs

Error message

psend is currently only implemented on GPUs

What it means

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.

Source

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

batching.fancy_primitive_batchers[ppermute_p] = _ppermute_batcher


@dataclass(frozen=True, slots=True)
class SingleSideCollectiveEffect(core.Effect):
  __str__ = lambda _: "one-sided communication"
  def __hash__(self):
    return hash(SingleSideCollectiveEffect)
  def __eq__(self, other):
    return isinstance(other, SingleSideCollectiveEffect)


single_side_collective_effect = SingleSideCollectiveEffect()
core.effects.control_flow_allowed_effects.add_type(SingleSideCollectiveEffect)

def _psend_lowering_gpu(ctx, x, *, axis_name, perm):
  if ("cuda" not in ctx.module_context.platforms and
      "rocm" not in ctx.module_context.platforms):
    raise NotImplementedError("psend is currently only implemented on GPUs")

  full_perm, other_args = _pcollectives_lowering_common(
      ctx, axis_name=axis_name, perm=perm, op_name="psend"
  )
  token = hlo.create_token()
  send_op = hlo.SendOp(
      [x],
      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):
    raise NotImplementedError("psend currently only supports manual sharding")

  return send_op.results

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Run on a CUDA or ROCm backend (set jax.default_device to a GPU or run on GPU hosts)
2. Guard psend/precv usage with jax.devices() platform checks and provide a CPU fallback
3. Request/await CPU support upstream

Example fix

// before
y = jax.jit(lambda x, t: lax.psend(x, t, 'i', perm))(x, token)
// after
if jax.default_backend() in ('gpu', 'cuda', 'rocm'):
    y = jax.jit(lambda x, t: lax.psend(x, t, 'i', perm))(x, token)
else:
    y = cpu_fallback(x)

import jax
assert jax.default_backend() in ('gpu', 'cuda', 'rocm'), 'psend requires a GPU backend'

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

try:
    y = jax.jit(f)(x)
except NotImplementedError as e:
    if 'psend' in str(e): y = cpu_stub(x)
    else: raise

When it happens

Trigger: Compiling a function containing jax.lax.psend on CPU or TPU (jax.default_device or available backends).

Common situations: Developing/testing distributed code locally on CPU before running on GPU clusters; CI without GPUs hitting the psend path.

Related errors

• psend is currently only implemented on GPU
• __dlpack__ device only supported for CPU, GPU and TPU pinned
• `buffer_callback` not supported on {platform} backend.
• the platform for the specified backend {xb.canonicalize_plat
• QDWH implementation is only supported on TPU

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