{"record":{"id":"b78112c4846a3d96","repo":"jax-ml/jax","slug":"explicit-sharding-inference-for-ragged-dot-general","errorCode":null,"errorMessage":"Explicit sharding inference for ragged_dot_general is not currently implemented. Please specify out_sharding.","messagePattern":"Explicit sharding inference for ragged_dot_general is not currently implemented\\. Please specify out_sharding\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6700,"sourceCode":"  if _is_ragged_contracting(batched_args[0].ndim - 1,\n                            ragged_dot_dimension_numbers):\n    result_batch_dim += 1\n  return batched_out, result_batch_dim\n\n\ndef _ragged_dot_general_sharding_rule(\n    lhs, rhs, group_sizes, *, ragged_dot_dimension_numbers, precision,\n    preferred_element_type: DTypeLike | None, group_offset, out_sharding):\n  mesh_set = {x.sharding.mesh for x in [lhs, rhs, group_sizes]\n              if not x.sharding.mesh.empty}\n  if len(mesh_set) > 1:\n    raise core.ShardingTypeError(\n      'All argument meshes must be the same or unspecified, but got'\n      f' lhs mesh = {lhs.sharding.mesh}, rhs mesh = {rhs.sharding.mesh},'\n      f' group_sizes mesh = {group_sizes.sharding.mesh}')\n\n  if out_sharding is None:\n    raise NotImplementedError(\n      \"Explicit sharding inference for ragged_dot_general is not currently\"\n      \" implemented. Please specify out_sharding.\")\n  return out_sharding\n\n\nragged_dot_general_p = standard_primitive(\n    _ragged_dot_general_shape_rule,\n    _ragged_dot_general_dtype_rule,\n    'ragged_dot_general',\n    vma_rule=partial(core.standard_vma_rule, 'ragged_dot'),\n    sharding_rule=_ragged_dot_general_sharding_rule,\n)\nad.primitive_jvps[ragged_dot_general_p] = _ragged_dot_general_jvp_rule\nad.primitive_transposes[ragged_dot_general_p] = _ragged_dot_general_transpose_rule\nbatching.fancy_primitive_batchers[ragged_dot_general_p] = _ragged_dot_general_batch_rule\n\n\ndef _ragged_dot_general_impl(","sourceCodeStart":6682,"sourceCodeEnd":6718,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6682-L6718","documentation":"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.","triggerScenarios":"Calling ragged_dot_general on sharded (jax.Array with NamedSharding) operands without passing out_sharding, under SPMD/multi-device execution.","commonSituations":"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.","solutions":["Pass out_sharding explicitly, e.g. NamedSharding(mesh, P('batch','model')) or a PartitionSpec","Compute the expected output sharding from the lhs sharding minus the ragged dimension's partitioning","Run unsharded (device_put to a single device) if sharding isn't required"],"exampleFix":"// before\nout = ragged_dot_general(x, w, gs, dn, mode=mode)  # x,w sharded\n// after\nfrom jax.sharding import NamedSharding, PartitionSpec as P\nout = ragged_dot_general(x, w, gs, dn, mode=mode,\n                         out_sharding=NamedSharding(mesh, P('batch', None)))","handlingStrategy":"validation","validationCode":"from jax import device_get\nif isinstance(x, jax.Array) and not isinstance(x.sharding, jax.sharding.SingleDeviceSharding):\n    assert out_sharding is not None","typeGuard":"def needs_out_sharding(*arrays) -> bool:\n    return any(len(getattr(a, 'sharding', SingleDeviceSharding()).device_set) > 1 for a in arrays)","tryCatchPattern":"except NotImplementedError as e:\n    if 'out_sharding' in str(e): retry with explicit NamedSharding(mesh, P(...))","preventionTips":["Always pass out_sharding when operands are sharded","Compute output spec alongside input specs in one place"],"tags":["jax","ragged-dot","sharding","multi-device","spmd"],"backgroundTag":"missing-sharding-specification","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}