{"record":{"id":"39711c3eb35228b6","repo":"jax-ml/jax","slug":"unsupported-memory-space-orig-ref-ty-memory-spac","errorCode":null,"errorMessage":"Unsupported memory space: {orig_ref_ty.memory_space}","messagePattern":"Unsupported memory space: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":468,"sourceCode":"        vec_size=out_layout.vec_size,\n    )\n    return [_fragmented_array_to_ir(fragmented_array)]\n\n  if not isinstance(out_layout, fa.TiledLayout):\n    raise ValueError(f\"{op} has an unsupported layout: {out_layout_attr}\")\n\n  optimized = op.optimized.value if op.optimized is not None else None\n  if transformed_ref.type.memory_space is None:  # GMEM\n    fragmented_array = fa.FragmentedArray.load_untiled(\n        transformed_ref,\n        layout=out_layout,\n        is_signed=is_signed,\n        optimized=bool(optimized),\n    )\n    return [_fragmented_array_to_ir(fragmented_array)]\n\n  if transforms_attr is None:\n    raise ValueError(f\"Unsupported memory space: {orig_ref_ty.memory_space}\")\n\n  swizzle = swizzle_from_transforms_attr(transforms_attr)\n  transforms = memref_transforms_from_transforms_attr(transforms_attr)\n  if transforms:\n    [tiling_transform] = transforms\n    assert isinstance(tiling_transform, lc.TileTransform)\n\n    def load_tiled(optimized: bool) -> fa.FragmentedArray:\n      return fa.FragmentedArray.load_tiled(\n          transformed_ref,\n          swizzle,\n          is_signed=is_signed,\n          layout=out_layout,\n          optimized=optimized,\n          tiling_rank=len(tiling_transform.tiling)\n      )\n\n    fragmented_array = _retry_on_failure(load_tiled, optimized)","sourceCodeStart":450,"sourceCodeEnd":486,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L450-L486","documentation":"The vector load lowering hit a reference whose memory space is neither SMEM nor plain GMEM-with-transforms, i.e. transforms_attr is None for a memory space it cannot handle. Only GMEM with an attached transforms attribute (tiling/swizzle) is supported on this path.","triggerScenarios":"Loading from a memref whose memory_space is set (e.g. #mgpu.thread or another address space) or a GMEM ref without a transforms attribute, hitting _vector_load_op_lowering_rule's fallback branch.","commonSituations":"Loading directly from a reference to special address spaces (e.g. tensor-memory or thread-local memory) instead of SMEM; constructing refs manually via memref ops bypassing Mosaic helpers.","solutions":["Load from shared memory or a plain GMEM ref created by mosaic's alloc helpers","Ensure the ref carries in_transforms (tiling/swizzle) if loading from GMEM","Avoid mgpu.memref_load style ops on exotic memory spaces; use mosaic's smem utilities"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"space = ref.type.memory_space\nassert space is None or int(space) == 1 or transforms_attr is not None, f'unsupported memory space {space}'","typeGuard":"def loadable_ref(ref, transforms_attr) -> bool:\n    return ref.type.memory_space is None and transforms_attr is not None","tryCatchPattern":null,"preventionTips":["Only load from GMEM refs with in_transforms or from SMEM","Use mosaic alloc helpers rather than custom memrefs"],"tags":["jax","mosaic","gpu","memory-space","memref"],"backgroundTag":"unsupported-memory-space","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}