jax-ml/jax · error · NotImplementedError

Only slicing with static indices allowed.

Error message

Only slicing with static indices allowed.

What it means

vector.extract in Mosaic lowering only supports static (compile-time constant) positions; dynamic index extraction is not implemented.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:794

  a = _fragmented_array_from_ir(op.source, in_layout)
  indices = tuple(
      utils.DynamicSlice(
          ir.IntegerAttr(offset).value, ir.IntegerAttr(length).value
      )
      for offset, length in zip(op.offsets, op.sizes, strict=True)
  )
  result = a[indices]
  assert result.layout == layouts_lib.from_layout_attr(out_layout)
  return [fragmented_array_to_ir(result, out_vec_ty)]


@_register_lowering(vector.ExtractOp)
def _vector_extract_op_lowering_rule(
    ctx: LoweringContext, op: vector.ExtractOp
) -> Sequence[ir.Value]:
  del ctx
  if op.dynamic_position:
    raise NotImplementedError("Only slicing with static indices allowed.")

  [in_layout] = inference_utils.in_layouts(op)
  a = _fragmented_array_from_ir(op.source, in_layout)

  if not isinstance(op.result.type, ir.VectorType):  # scalar result
    result = a[tuple(op.static_position)]
    assert isinstance(result.layout, fa.WGSplatFragLayout)
    return [result.registers.item()]

  [out_layout] = inference_utils.out_layouts(op)
  assert in_layout == out_layout
  a = _fragmented_array_from_ir(op.source, in_layout)
  result_type = ir.VectorType(op.result.type)
  result = a[tuple(op.static_position)]
  assert result.layout == layouts_lib.from_layout_attr(out_layout)
  return [fragmented_array_to_ir(result, result_type)]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the extraction index a Python/compile-time constant
2. Store the vector to memory and load the dynamic element from SMEM/GMEM instead
3. Unroll loops in Python so indices become static

Example fix

// before
i = arith.index  # dynamic
vector.extract v[i]
// after
for i in range(static_n):  # unrolled, static
    vector.extract v[i]

assert not op.dynamic_position, 'extract indices must be static constants'

When it happens

Trigger: Emitting vector.extract whose dynamic_position is non-empty (indices supplied as runtime values).

Common situations: Trying to index a vector register by a loop variable computed at runtime inside the kernel.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• `strides` must contain only 1s.
• Unsupported reduction kind: {op.kind}
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}
• Unsupported memory space: {orig_ref_ty.memory_space}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/275f5cf3b31eaade. Report an issue: GitHub.