jax-ml/jax · error · NotImplementedError

precv currently only supports manual sharding

Error message

precv currently only supports manual sharding

What it means

precv's GPU lowering requires an SPMDAxisContext: the call must be under jax.shard_map with manual sharding so that receiver replica semantics are defined. Other axis contexts (automatic sharding, eager) are rejected.

Source

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

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):
    raise NotImplementedError("precv currently only supports manual sharding")

  # 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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap precv (and matching psend) in jax.shard_map over a mesh
2. Give inputs a NamedSharding matching the mesh
3. Verify with jax.debug that the axis context is SPMD before calling

Example fix

// before
y = lax.precv(token, shape, 'i', perm)
// after
y = jax.shard_map(lambda tok: lax.precv(tok, shape, 'i', perm), mesh)(token)

null  # structural: ensure precv appears only inside shard_map bodies

When it happens

Trigger: Calling jax.lax.precv outside jax.shard_map, or under jit with automatic (GSD) sharding.

Common situations: Building pipeline-parallel send/recv pairs and forgetting the shard_map wrapper on the receiving side.

Related errors

• psend currently only supports manual sharding
• Mesh must be provided for shard_map with checkify.
• Unsupported aval type: {type(v)}
• Primitive {prim_name} requires varying manual axes to match,
• {} function carry input and carry output must have equal typ

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