{"record":{"id":"63db205c7c79e6a4","repo":"jax-ml/jax","slug":"only-leading-gather-dimensions-allowed","errorCode":null,"errorMessage":"Only leading gather dimensions allowed.","messagePattern":"Only leading gather dimensions allowed\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2222,"sourceCode":"\n@_add_constraint_system_derivation_rule(mgpu.AsyncLoadOp)\n@_add_constraint_system_derivation_rule(mgpu.AsyncStoreOp)\ndef _async_load_store_constraint_system(\n    ctx: DerivationContext,\n    op: mgpu.AsyncLoadOp | mgpu.AsyncStoreOp,\n) -> ConstraintSystemDerivationRuleResult:\n  # We only support 2D gathers/scatters along the leading dimension. Tiling\n  # either keeps the gather/scatter dimension leading or allows\n  # collapsing leading dimensions to maintain contiguity without\n  # transforming global memory.\n  tiling_multiple = []\n  for i, (size, index) in enumerate(zip(op.slice_lengths, op.indices, strict=True)):\n    if size == -1:\n      # This dimension does not appear in the final smem memref shape.\n      continue\n    if isinstance(index.type, ir.VectorType):\n      if i != 0:\n        raise NotImplementedError(\"Only leading gather dimensions allowed.\")\n      if isinstance(op, mgpu.AsyncStoreOp):\n        gmem_shape = ir.MemRefType(op.destination.type).shape\n      else:\n        gmem_shape = ir.MemRefType(op.source.type).shape\n      if len(gmem_shape) != 2:\n        raise NotImplementedError(\"Only 2D gathers/scatters for async load/store are supported.\")\n      tiling_multiple.append(size)\n      continue\n    tiling_multiple.append(dynamic_gcd(size, index))\n\n  operand_index = 1 if isinstance(op, mgpu.AsyncLoadOp) else 0\n  operand = ValueSite(op, VariableType.OPERAND, operand_index)\n  var = ctx.producer_ref(operand)\n  constraints: list[cs.Constraint] = [\n      cs.Divides(expr=var, tiling_multiple=tuple(tiling_multiple))\n  ]\n  if any(isinstance(idx.type, ir.VectorType) for idx in op.indices):\n    element_bitwidth = utils.bitwidth(op.source.type.element_type)","sourceCodeStart":2204,"sourceCodeEnd":2240,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2204-L2240","documentation":"In async gather/scatter (vector-indexed async load/store), Mosaic only supports gather dimensions that appear first: if a vector-typed index occurs at any position other than i == 0, layout inference raises NotImplementedError, because mixed leading-static/gather dims would need more general constraint modeling.","triggerScenarios":"Constructing mgpu.async_load or mgpu.async_store where slice_lengths[i] != -1 with a VectorType indices[i] at position i > 0 — e.g. gathering along the second (column) dimension instead of the first (row) dimension of a 2D tensor.","commonSituations":"Column gathers (indexing the fast/last axis with vectors), or building indices tuples where a static scalar index precedes the vector index; porting Triton gather code that gathers along axis 1.","solutions":["Transpose the global-memory tensor (or its layout) so the gathered dimension becomes dimension 0, then gather along rows","Reorder the indices tuple so the vector index is the first operand and preceding dims use size == -1 (sliced-away dims)","Use non-async per-element load/store for non-leading gather dims if performance permits"],"exampleFix":"# before: gather along dim 1 (i == 1 with VectorType) -> NotImplementedError\nmgpu.async_load(src_T, smem, indices=(scalar_row, vec_cols), slice_lengths=(1, 1))\n\n# after: transpose so gather is along dim 0\nsrc = transpose(src_T)\nmgpu.async_load(src, smem, indices=(vec_cols, scalar_row_or_none), slice_lengths=(1, 1))","handlingStrategy":"validation","validationCode":"for i, (size, idx) in enumerate(zip(slice_lengths, indices)):\n    if size != -1 and not isinstance(idx.type, ir.VectorType):\n        continue\n    if size != -1 and isinstance(idx.type, ir.VectorType) and i != 0:\n        raise ValueError('transpose tensors so gather dim is dim 0')","typeGuard":"def gather_dims_are_leading(slice_lengths, indices) -> bool:\n    return all(\n        size == -1 or not isinstance(idx.type, ir.VectorType) or i == 0\n        for i, (size, idx) in enumerate(zip(slice_lengths, indices))\n    )","tryCatchPattern":null,"preventionTips":["Gather only along dimension 0; transpose gmem layout otherwise","Keep indices tuples ordered: gather (vector) index first"],"tags":["jax","mosaic-gpu","tma","gather","async-load","layout-inference"],"backgroundTag":"unsupported-operation-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}