{"record":{"id":"055b253e513e977a","repo":"jax-ml/jax","slug":"cannot-retrieve-the-architecture-no-module-found","errorCode":null,"errorMessage":"Cannot retrieve the architecture: no module found","messagePattern":"Cannot retrieve the architecture: no module found","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":2449,"sourceCode":"\n\ndef get_arch() -> Arch:\n  ip = ir.InsertionPoint.current\n  if ip is None:\n    raise ValueError(\n        \"Cannot retrieve the architecture without an insertion point\"\n    )\n  block = ip.block\n  op = block.owner\n  while op is not None:\n    if op.name == \"builtin.module\":\n      arch_major = op.attributes[\"mosaic_gpu.arch_major\"]\n      arch_minor = op.attributes[\"mosaic_gpu.arch_minor\"]\n      assert isinstance(arch_major, ir.IntegerAttr)\n      assert isinstance(arch_minor, ir.IntegerAttr)\n      return Arch(arch_major.value, arch_minor.value)\n    op = op.parent\n  raise ValueError(\"Cannot retrieve the architecture: no module found\")\n\n\ndef reduce_shape(\n    shape: Sequence[int], axes: Sequence[int], keep_dims: bool = False\n) -> tuple[int, ...]:\n  res = []\n  for i, dim in enumerate(shape):\n    if i in axes:\n      if keep_dims:\n        res.append(1)\n    else:\n      res.append(dim)\n  return tuple(res)\n","sourceCodeStart":2431,"sourceCodeEnd":2463,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L2431-L2463","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Run the code through the mosaic_gpu pipeline (which stamps arch attributes on the module) rather than manual IR building","Set the attributes manually: module.operation.attributes['mosaic_gpu.arch_major'] = ir.IntegerAttr.get(ir.IndexType.get(), 90)","Upgrade JAX if your version predates automatic arch stamping","Pass arch explicitly through your code instead of relying on get_arch"],"exampleFix":"# before\nwith ir.Context(), ir.Location.unknown():\n    mod = ir.Module.create()\n    with ir.InsertionPoint(mod.body):\n        arch = utils.get_arch()  # ValueError\n# after\nmod.operation.attributes['mosaic_gpu.arch_major'] = ir.IntegerAttr.get(ir.i32(), 90)\nmod.operation.attributes['mosaic_gpu.arch_minor'] = ir.IntegerAttr.get(ir.i32(), 0)\nwith ir.InsertionPoint(mod.body):\n    arch = utils.get_arch()","handlingStrategy":"validation","validationCode":"mod.operation.attributes.setdefault('mosaic_gpu.arch_major', ir.IntegerAttr.get(ir.i32(), 90))\nmod.operation.attributes.setdefault('mosaic_gpu.arch_minor', ir.IntegerAttr.get(ir.i32(), 0))","typeGuard":null,"tryCatchPattern":"try:\n    arch = utils.get_arch()\nexcept ValueError as e:\n    if 'no module found' in str(e):\n        raise RuntimeError('run through mosaic_gpu pipeline or stamp arch attrs') from e\n    raise","preventionTips":["In tests, stamp arch_major/arch_minor on hand-built modules","Run kernels through the mosaic_gpu compilation entry points"],"tags":["jax","mosaic-gpu","architecture","mlir-module"],"backgroundTag":"mlir-context-missing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}