{"record":{"id":"fabd74560de880c2","repo":"jax-ml/jax","slug":"explicit-sharding-is-not-currently-supported-in-th","errorCode":null,"errorMessage":"Explicit sharding is not currently supported in the pallas-triton lowering of ragged_dot_general. You can call this op under shard_map.","messagePattern":"Explicit sharding is not currently supported in the pallas-triton lowering of ragged_dot_general\\. You can call this op under shard_map\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py","lineNumber":462,"sourceCode":"  ds = list(backend.devices())\n  if not ds:\n    return False\n  return tuple(int(x) for x in ds[0].compute_capability.split(\".\")) >= (8, 0)\n\n\ndef _pallas_ragged_dot_general_impl(\n    lhs: Array,\n    rhs: Array,\n    group_sizes: Array,\n    ragged_dot_dimension_numbers: lax.RaggedDotDimensionNumbers,\n    precision: lax.PrecisionLike = None,  # TODO(rdyro): Add support for the\n                                          # precision parameter in the kernels.\n    preferred_element_type: DTypeLike | None = None,\n    group_offset: Array | None = None,\n    out_sharding=None,\n) -> Array:\n  if out_sharding is not None:\n    raise NotImplementedError(\n      \"Explicit sharding is not currently supported in the pallas-triton\"\n      \" lowering of ragged_dot_general. You can call this op under shard_map.\")\n\n  if group_offset is not None:\n    raise NotImplementedError(\"group_offset is not currently supported in the \"\n                              \"pallas-triton lowering.\")\n\n  mode, lhs_ragged_dim = lax._ragged_dot_mode_and_dim(\n      lhs.ndim, ragged_dot_dimension_numbers)\n  (l_contract, r_contract), (l_batch, r_batch) = (\n      ragged_dot_dimension_numbers.dot_dimension_numbers\n  )\n  l_contract, r_contract = tuple(l_contract), tuple(r_contract)\n  l_batch, r_batch = tuple(l_batch), tuple(r_batch)\n  l_noncontract = tuple(lax.remaining(range(lhs.ndim), l_contract, l_batch))\n  assert len(l_batch) == len(r_batch)\n\n  if group_sizes.ndim == 1:","sourceCodeStart":444,"sourceCodeEnd":480,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/pallas_lowerings/gpu/ragged_dot.py#L444-L480","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Wrap the ragged_dot_general call in jax.experimental.shard_map.shard_map and let sharding be implied by the in_shardings","Drop out_sharding (pass None) and reshard afterwards with jax.lax.with_sharding_constraint","Switch to the non-Pallas lowering (lax.ragged_dot_general direct XLA path) if explicit sharding is essential"],"exampleFix":"// before\nout = lax.ragged_dot_general(lhs, rhs, ..., out_sharding=sharding)\n\n// after\nout = shard_map(lambda l, r: lax.ragged_dot_general(l, r, ...), mesh,\n                in_specs=(spec_l, spec_r), out_specs=spec_out)(lhs, rhs)","handlingStrategy":"fallback","validationCode":"assert out_sharding is None or using_xla_path, 'use shard_map instead of out_sharding'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer shard_map for SPMD Pallas kernels","Guard backend-specific args behind feature flags per backend"],"tags":["jax","sharding","pallas","spmd","not-implemented"],"backgroundTag":"feature-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}