jax-ml/jax · error · NotImplementedError

psend currently only supports manual sharding

Error message

psend currently only supports manual sharding

What it means

psend's lowering requires an SPMDAxisContext, i.e. the code must be under jax.shard_map with manual sharding so replica semantics apply. Under automatic (GSD) sharding or eager execution the axis context is a different type and lowering is refused.

Source

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

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


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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the psend/precv code in jax.shard_map with mesh and manual sharding
2. Ensure the sharded inputs have a sharding compatible with the mesh (e.g. NamedSharding(mesh, ...))
3. Move communication logic into a shard_map function passed to jit

Example fix

// before
y = jax.jit(lambda x, t: lax.psend(x, t, 'i', perm))(x, token)
// after
y = jax.jit(jax.shard_map(lambda x, t: lax.psend(x, t, 'i', perm), mesh))(x, token)

assert in_shard_map, 'psend must be called inside jax.shard_map (manual sharding)'  # enforce via code review/lint

null

When it happens

Trigger: Calling jax.lax.psend outside jax.shard_map (e.g. jit with automatic sharding, or traced eagerly).

Common situations: Writing point-to-point communication pipelines and forgetting the shard_map wrapper; assuming jit+NamedSharding provides manual sharding.

Related errors

• precv 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/519f29f51aca4e7b. Report an issue: GitHub.