jax-ml/jax · error · NotImplementedError

Explicit sharding is not currently supported in the pallas-t

Error message

Explicit sharding is not currently supported in the pallas-triton lowering of ragged_dot_general. You can call this op under shard_map.

What it means

The pallas-triton lowering of ragged_dot_general does not implement explicit output sharding (the out_sharding argument). If you need the result distributed across devices, you must express the op inside shard_map so sharding is inferred contextually.

Source

Thrown at jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py:462

  ds = list(backend.devices())
  if not ds:
    return False
  return tuple(int(x) for x in ds[0].compute_capability.split(".")) >= (8, 0)


def _pallas_ragged_dot_general_impl(
    lhs: Array,
    rhs: Array,
    group_sizes: Array,
    ragged_dot_dimension_numbers: lax.RaggedDotDimensionNumbers,
    precision: lax.PrecisionLike = None,  # TODO(rdyro): Add support for the
                                          # precision parameter in the kernels.
    preferred_element_type: DTypeLike | None = None,
    group_offset: Array | None = None,
    out_sharding=None,
) -> Array:
  if out_sharding is not None:
    raise NotImplementedError(
      "Explicit sharding is not currently supported in the pallas-triton"
      " lowering of ragged_dot_general. You can call this op under shard_map.")

  if group_offset is not None:
    raise NotImplementedError("group_offset is not currently supported in the "
                              "pallas-triton lowering.")

  mode, lhs_ragged_dim = lax._ragged_dot_mode_and_dim(
      lhs.ndim, ragged_dot_dimension_numbers)
  (l_contract, r_contract), (l_batch, r_batch) = (
      ragged_dot_dimension_numbers.dot_dimension_numbers
  )
  l_contract, r_contract = tuple(l_contract), tuple(r_contract)
  l_batch, r_batch = tuple(l_batch), tuple(r_batch)
  l_noncontract = tuple(lax.remaining(range(lhs.ndim), l_contract, l_batch))
  assert len(l_batch) == len(r_batch)

  if group_sizes.ndim == 1:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the ragged_dot_general call in jax.experimental.shard_map.shard_map and let sharding be implied by the in_shardings
2. Drop out_sharding (pass None) and reshard afterwards with jax.lax.with_sharding_constraint
3. Switch to the non-Pallas lowering (lax.ragged_dot_general direct XLA path) if explicit sharding is essential

Example fix

// before
out = lax.ragged_dot_general(lhs, rhs, ..., out_sharding=sharding)

// after
out = shard_map(lambda l, r: lax.ragged_dot_general(l, r, ...), mesh,
                in_specs=(spec_l, spec_r), out_specs=spec_out)(lhs, rhs)

assert out_sharding is None or using_xla_path, 'use shard_map instead of out_sharding'

When it happens

Trigger: Passing out_sharding=... to ragged_dot_general when the pallas-triton lowering is selected (e.g. inside a sharded MoE path on GPU/triton backend).

Common situations: Migrating code from the TPU ragged_dot path (which supports out_sharding) to GPU; adding SPMD annotations to an MoE kernel during multi-GPU scaling.

Related errors

• pallas_call with a mesh does not support batching
• callbacks are only supported in spmd computations when all m
• callbacks do not support specifying sharding inside spmd com
• {name} cannot accept args which are unreduced. Got {a.str_sh
• {name} cannot accept args with unreduced_kind={a.mat.unreduc

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