{"record":{"id":"275f5cf3b31eaade","repo":"jax-ml/jax","slug":"only-slicing-with-static-indices-allowed","errorCode":null,"errorMessage":"Only slicing with static indices allowed.","messagePattern":"Only slicing with static indices allowed\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":794,"sourceCode":"  a = _fragmented_array_from_ir(op.source, in_layout)\n  indices = tuple(\n      utils.DynamicSlice(\n          ir.IntegerAttr(offset).value, ir.IntegerAttr(length).value\n      )\n      for offset, length in zip(op.offsets, op.sizes, strict=True)\n  )\n  result = a[indices]\n  assert result.layout == layouts_lib.from_layout_attr(out_layout)\n  return [fragmented_array_to_ir(result, out_vec_ty)]\n\n\n@_register_lowering(vector.ExtractOp)\ndef _vector_extract_op_lowering_rule(\n    ctx: LoweringContext, op: vector.ExtractOp\n) -> Sequence[ir.Value]:\n  del ctx\n  if op.dynamic_position:\n    raise NotImplementedError(\"Only slicing with static indices allowed.\")\n\n  [in_layout] = inference_utils.in_layouts(op)\n  a = _fragmented_array_from_ir(op.source, in_layout)\n\n  if not isinstance(op.result.type, ir.VectorType):  # scalar result\n    result = a[tuple(op.static_position)]\n    assert isinstance(result.layout, fa.WGSplatFragLayout)\n    return [result.registers.item()]\n\n  [out_layout] = inference_utils.out_layouts(op)\n  assert in_layout == out_layout\n  a = _fragmented_array_from_ir(op.source, in_layout)\n  result_type = ir.VectorType(op.result.type)\n  result = a[tuple(op.static_position)]\n  assert result.layout == layouts_lib.from_layout_attr(out_layout)\n  return [fragmented_array_to_ir(result, result_type)]\n\n","sourceCodeStart":776,"sourceCodeEnd":812,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L776-L812","documentation":"vector.extract in Mosaic lowering only supports static (compile-time constant) positions; dynamic index extraction is not implemented.","triggerScenarios":"Emitting vector.extract whose dynamic_position is non-empty (indices supplied as runtime values).","commonSituations":"Trying to index a vector register by a loop variable computed at runtime inside the kernel.","solutions":["Make the extraction index a Python/compile-time constant","Store the vector to memory and load the dynamic element from SMEM/GMEM instead","Unroll loops in Python so indices become static"],"exampleFix":"// before\ni = arith.index  # dynamic\nvector.extract v[i]\n// after\nfor i in range(static_n):  # unrolled, static\n    vector.extract v[i]","handlingStrategy":"validation","validationCode":"assert not op.dynamic_position, 'extract indices must be static constants'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Unroll Python loops so extract indices are compile-time constants","Round-trip through SMEM for dynamic indexing"],"tags":["jax","mosaic","gpu","vector","static-index"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}