jax-ml/jax · error · NotImplementedError

precv is currently only implemented on GPU

Error message

precv is currently only implemented on GPU

What it means

Fallback lowering for precv_p on non-GPU platforms: precv only has a GPU (CUDA/ROCm) lowering rule, so compiling it elsewhere hits this unconditional NotImplementedError with a clear message.

Source

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

  # recv_op should return an array of [RankedTensorType, StableHlo.token]; we
  # only need the tensor.
  results = recv_op.results
  return [results[0]]


def _precv_abstract_eval(
    token, *, out_shape, axis_name, **params
):
  return out_shape, {*map(core.NamedAxisEffect, axis_name),
                     single_side_collective_effect}

precv_p = core.Primitive("precv")
precv_p.def_effectful_abstract_eval(_precv_abstract_eval)
mlir.register_lowering(precv_p, _precv_lowering_gpu, platform='gpu')

def _precv_lowering(ctx, token, *, out_shape, axis_name, perm):
  raise NotImplementedError("precv is currently only implemented on GPU")
mlir.register_lowering(precv_p, _precv_lowering)

batching.fancy_primitive_batchers[precv_p] = _ppermute_batcher

def _pbroadcast_transpose_rule(t, x, source, axis_name):
  is_source = axis_index(axis_name) == source
  tsum = psum(t, axis_name)
  return [lax.select(is_source, lax.full_like(t, tsum), lax.full_like(t, 0))]

def _pbroadcast_batcher(axis_data, vals_in, dims_in, axis_name, source):
  axis_size = axis_data.size
  (v,), (d,) = vals_in, dims_in
  if not isinstance(axis_name, (tuple, list)):
    axis_name = (axis_name,)
  if d is None and axis_data.name not in axis_name:
    return pbroadcast_p.bind(v, axis_name=axis_name, source=source), None
  if axis_data.name not in axis_name:
    return pbroadcast_p.bind(v, axis_name=axis_name, source=source), d

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Run on CUDA/ROCm backends
  2. Gate precv paths on backend capability
  3. Provide a non-NCCL fallback (e.g. all_gather based) for CPU tests

Example fix

// before
y = jax.jit(f_with_precv)(token)
// after
if jax.default_backend() != 'gpu':
    y = mock_recv_fallback(token)
else:
    y = jax.jit(f_with_precv)(token)
Defensive patterns

Strategy: type-guard

Validate before calling

if jax.default_backend() not in ('gpu', 'cuda', 'rocm'):
    pytest.skip('precv requires GPU')

Type guard

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

Try / catch

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

Prevention

When it happens

Trigger: Compiling code containing lax.precv on CPU or TPU.

Common situations: CPU local testing of distributed pipelines; GPU-less CI.

Related errors

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