{"record":{"id":"88444ff326282d6e","repo":"jax-ml/jax","slug":"unsupported-tma-index-shape-shape","errorCode":null,"errorMessage":"Unsupported TMA index shape {shape}","messagePattern":"Unsupported TMA index shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2199,"sourceCode":"      base_operand_index = 3\n    case mgpu.AsyncStoreOp():\n      base_operand_index = 2\n    case mgpu.AsyncPrefetchOp():\n      base_operand_index = 1\n\n  for i, idx in enumerate(op.indices):\n    if isinstance(idx.type, ir.VectorType):\n      value_site = ValueSite(op, VariableType.OPERAND, base_operand_index + i)\n      value_site_var = cs.Variable(value_site)\n      shape = tuple(idx.type.shape)\n\n      allowed_layouts = []\n      if shape[0] % 16 == 0:\n        allowed_layouts.append(cs.RegisterLayout(value=fa.TMA_INDICES_LAYOUT))\n      if shape[0] % 4 == 0:\n        allowed_layouts.append(cs.RegisterLayout(value=fa.TMA_INDICES_4_LAYOUT))\n      if not allowed_layouts:\n        raise ValueError(f\"Unsupported TMA index shape {shape}\")\n      values_sites[value_site_var] = [value_site]\n      constraints.append(cs.OneOf(value_site_var, tuple(allowed_layouts)))\n  return values_sites, constraints\n\n\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:","sourceCodeStart":2181,"sourceCodeEnd":2217,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2181-L2217","documentation":"For async TMA load/store/prefetch ops, vector-valued indices must have a first dimension divisible by 16 (TMA_INDICES_LAYOUT) or by 4 (TMA_INDICES_4_LAYOUT). If neither divides shape[0], no register layout exists for the index vector and a ValueError is raised.","triggerScenarios":"Passing an index vector whose leading dimension is not a multiple of 4 — e.g. shape[0] == 6 or 10 — as the indices operand of mgpu.async_load/async_store/async_prefetch in gather/scatter mode.","commonSituations":"Gather kernels with an arbitrary number of rows gathered per program; changing the gather batch size to a value not divisible by 4/16 after tuning; padding index vectors inconsistently with the TMA hardware constraints.","solutions":["Pad or reshape the index vector so its first dimension is a multiple of 4 (or 16) — pad with dummy/unused indices and mask them out","Choose gather batch sizes that are multiples of 16 to get the more efficient TMA_INDICES_LAYOUT","Split the gather into chunks whose index counts are multiples of 4/16"],"exampleFix":"# before: idx shape (6, ...) -> ValueError\nmgpu.async_load(source, smem, indices=idx6, slice_lengths=(1, w), ...)\n\n# after: pad to multiple of 4 (mask or ignore extra rows)\nidx8 = pad_to_multiple(idx6, multiple=4)\nmgpu.async_load(source, smem, indices=idx8, slice_lengths=(1, w), ...)","handlingStrategy":"validation","validationCode":"n = indices.shape[0]\nif n % 16 and n % 4:\n    indices = pad_indices_to_multiple(indices, 4)  # and mask extras","typeGuard":"def is_valid_tma_index_shape(shape) -> bool:\n    return len(shape) >= 1 and (shape[0] % 16 == 0 or shape[0] % 4 == 0)","tryCatchPattern":null,"preventionTips":["Choose gather batch sizes that are multiples of 16","Pad index vectors to a multiple of 4 (ideally 16) and mask unused lanes"],"tags":["jax","mosaic-gpu","tma","async-load","indices","gather"],"backgroundTag":"shape-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}