jax-ml/jax · error · NotImplementedError

Explicit sharding inference for ragged_dot_general is not cu

Error message

Explicit sharding inference for ragged_dot_general is not currently implemented. Please specify out_sharding.

What it means

When running with explicit sharding (multi-device/mesh contexts), ragged_dot_general requires out_sharding to be provided; JAX cannot infer the output sharding for this primitive automatically.

Source

Thrown at jax/_src/lax/lax.py:6700

  if _is_ragged_contracting(batched_args[0].ndim - 1,
                            ragged_dot_dimension_numbers):
    result_batch_dim += 1
  return batched_out, result_batch_dim


def _ragged_dot_general_sharding_rule(
    lhs, rhs, group_sizes, *, ragged_dot_dimension_numbers, precision,
    preferred_element_type: DTypeLike | None, group_offset, out_sharding):
  mesh_set = {x.sharding.mesh for x in [lhs, rhs, group_sizes]
              if not x.sharding.mesh.empty}
  if len(mesh_set) > 1:
    raise core.ShardingTypeError(
      'All argument meshes must be the same or unspecified, but got'
      f' lhs mesh = {lhs.sharding.mesh}, rhs mesh = {rhs.sharding.mesh},'
      f' group_sizes mesh = {group_sizes.sharding.mesh}')

  if out_sharding is None:
    raise NotImplementedError(
      "Explicit sharding inference for ragged_dot_general is not currently"
      " implemented. Please specify out_sharding.")
  return out_sharding


ragged_dot_general_p = standard_primitive(
    _ragged_dot_general_shape_rule,
    _ragged_dot_general_dtype_rule,
    'ragged_dot_general',
    vma_rule=partial(core.standard_vma_rule, 'ragged_dot'),
    sharding_rule=_ragged_dot_general_sharding_rule,
)
ad.primitive_jvps[ragged_dot_general_p] = _ragged_dot_general_jvp_rule
ad.primitive_transposes[ragged_dot_general_p] = _ragged_dot_general_transpose_rule
batching.fancy_primitive_batchers[ragged_dot_general_p] = _ragged_dot_general_batch_rule


def _ragged_dot_general_impl(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass out_sharding explicitly, e.g. NamedSharding(mesh, P('batch','model')) or a PartitionSpec
2. Compute the expected output sharding from the lhs sharding minus the ragged dimension's partitioning
3. Run unsharded (device_put to a single device) if sharding isn't required

Example fix

// before
out = ragged_dot_general(x, w, gs, dn, mode=mode)  # x,w sharded
// after
from jax.sharding import NamedSharding, PartitionSpec as P
out = ragged_dot_general(x, w, gs, dn, mode=mode,
                         out_sharding=NamedSharding(mesh, P('batch', None)))

from jax import device_get
if isinstance(x, jax.Array) and not isinstance(x.sharding, jax.sharding.SingleDeviceSharding):
    assert out_sharding is not None

def needs_out_sharding(*arrays) -> bool:
    return any(len(getattr(a, 'sharding', SingleDeviceSharding()).device_set) > 1 for a in arrays)

except NotImplementedError as e:
    if 'out_sharding' in str(e): retry with explicit NamedSharding(mesh, P(...))

When it happens

Trigger: Calling ragged_dot_general on sharded (jax.Array with NamedSharding) operands without passing out_sharding, under SPMD/multi-device execution.

Common situations: Scaling a MoE model to multiple devices with jax.sharding NamedSharding on inputs; the same code works on CPU single-device but fails once arrays become sharded.

Related errors

• to_dlpack can only pack a dlpack tensor from an array on a s
• 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/b78112c4846a3d96. Report an issue: GitHub.