jax-ml/jax · error · ValueError

Cannot retrieve the architecture without an insertion point

Error message

Cannot retrieve the architecture without an insertion point

What it means

get_arch (utils.py:2436) determines the target GPU architecture by walking up from the current MLIR insertion point to the module and reading mosaic_gpu.arch_major/arch_minor attributes. Outside of an active insertion point there is no IR context to inspect, so it refuses.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:2436

  if element_bitwidth > 8:
    return arith.muli(offset, c(element_bitwidth // 8, index_ty))
  elif element_bitwidth < 8:
    return arith.divsi(offset, c(8 // element_bitwidth, index_ty))
  else:
    return offset


@dataclasses.dataclass(frozen=True)
class Arch:
  major: int
  minor: int


def get_arch() -> Arch:
  ip = ir.InsertionPoint.current
  if ip is None:
    raise ValueError(
        "Cannot retrieve the architecture without an insertion point"
    )
  block = ip.block
  op = block.owner
  while op is not None:
    if op.name == "builtin.module":
      arch_major = op.attributes["mosaic_gpu.arch_major"]
      arch_minor = op.attributes["mosaic_gpu.arch_minor"]
      assert isinstance(arch_major, ir.IntegerAttr)
      assert isinstance(arch_minor, ir.IntegerAttr)
      return Arch(arch_major.value, arch_minor.value)
    op = op.parent
  raise ValueError("Cannot retrieve the architecture: no module found")


def reduce_shape(
    shape: Sequence[int], axes: Sequence[int], keep_dims: bool = False
) -> tuple[int, ...]:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move the get_arch() call inside the kernel/lowering code where an insertion point is active
2. Wrap the call in 'with ir.InsertionPoint(module.body): ...' after entering an ir.Context
3. Pass Arch explicitly to your helper instead of inferring it
4. Use mosaic_gpu.current_arch() or the pipeline's arch plumbing if available in your version

Example fix

# before
ARCH = utils.get_arch()  # at module scope
# after
def kernel(...):
    arch = utils.get_arch()  # inside, insertion point active
    ...

assert ir.InsertionPoint.current is not None, 'call get_arch() inside a lowering'

try:
    arch = utils.get_arch()
except ValueError:
    arch = Arch(90, 0)  # fallback; must still stamp module attrs

When it happens

Trigger: Calling utils.get_arch() (directly or transitively, e.g. when constructing ops that depend on arch) from Python at import/tracing time with no active ir.InsertionPoint — no ongoing lowering or explicit 'with ir.InsertionPoint(block):' context.

Common situations: Building constants or computing shapes at module import time; calling in tests without setting up an MLIR context; helper code that runs before the mosaic lowering pipeline enters a function body.

Related errors

• Expected an index-typed index
• Cannot retrieve the architecture: no module found
• Sharding rule has {len(rule.operand_mappings)} operands, but
• Unsupported scalar attribute type: {type(val)}
• NumPy arrays with zero strides are not supported as MLIR att

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