jax-ml/jax · error · ValueError

Cannot retrieve the architecture: no module found

Error message

Cannot retrieve the architecture: no module found

What it means

get_arch (utils.py:2449) walked from the current insertion point up through parents and never found a 'builtin.module' op carrying mosaic_gpu.arch_major/arch_minor attributes. Without those attributes the target SM architecture cannot be determined.

Source

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

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, ...]:
  res = []
  for i, dim in enumerate(shape):
    if i in axes:
      if keep_dims:
        res.append(1)
    else:
      res.append(dim)
  return tuple(res)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Run the code through the mosaic_gpu pipeline (which stamps arch attributes on the module) rather than manual IR building
2. Set the attributes manually: module.operation.attributes['mosaic_gpu.arch_major'] = ir.IntegerAttr.get(ir.IndexType.get(), 90)
3. Upgrade JAX if your version predates automatic arch stamping
4. Pass arch explicitly through your code instead of relying on get_arch

Example fix

# before
with ir.Context(), ir.Location.unknown():
    mod = ir.Module.create()
    with ir.InsertionPoint(mod.body):
        arch = utils.get_arch()  # ValueError
# after
mod.operation.attributes['mosaic_gpu.arch_major'] = ir.IntegerAttr.get(ir.i32(), 90)
mod.operation.attributes['mosaic_gpu.arch_minor'] = ir.IntegerAttr.get(ir.i32(), 0)
with ir.InsertionPoint(mod.body):
    arch = utils.get_arch()

mod.operation.attributes.setdefault('mosaic_gpu.arch_major', ir.IntegerAttr.get(ir.i32(), 90))
mod.operation.attributes.setdefault('mosaic_gpu.arch_minor', ir.IntegerAttr.get(ir.i32(), 0))

try:
    arch = utils.get_arch()
except ValueError as e:
    if 'no module found' in str(e):
        raise RuntimeError('run through mosaic_gpu pipeline or stamp arch attrs') from e
    raise

When it happens

Trigger: Calling utils.get_arch() while building IR in a module created outside the mosaic_gpu pipeline (no arch attributes stamped), or in a manually constructed ir.Module in tests.

Common situations: Unit tests constructing MLIR modules by hand; older JAX/mosaic versions where attributes were not yet set; running lowering snippets outside the full mosaic_gpu compilation flow.

Related errors

• Cannot retrieve the architecture without an insertion point
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot specify both out_shape and out_type. Use out_type.
• Cannot specify both scratch_shapes and scratch_types. Use sc

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