{"record":{"id":"ab97a0aa236ebe1d","repo":"jax-ml/jax","slug":"unsupported-memory-space","errorCode":null,"errorMessage":"Unsupported memory space.","messagePattern":"Unsupported memory space\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":1670,"sourceCode":"          total_offset = base_offset + offset\n          ref_ty = ir.MemRefType.get(\n              transformed_shape, mlir_dtype, memory_space=mgpu_utils.tmem()\n          )\n          alloc_id = source_slice_op.alias_id\n          assert alloc_id is not None\n          # TODO(bchetioui): Use a scheme resilient to hash collisions.\n          alias_id = hash((offset, alloc_id.value, alias_group_idx))\n          slice_op = mgpu.dialect.SliceTmemOp(\n              ref_ty, source_slice_op.source, total_offset\n          )\n          i64 = ir.IntegerType.get_signless(64)\n          slice_op.attributes[\"alias_id\"] = ir.IntegerAttr.get(i64, alias_id)\n          ref = slice_op.result\n          assert layout is not None\n          layout_attr = mgpu.layouts.to_layout_attr(layout)\n          ref = mgpu.dialect.tmem_layout_cast(ref, layout_attr)\n        else:\n          raise NotImplementedError(\"Unsupported memory space.\")\n      return (\n          ref,\n          ref_aval,\n          transform_avals[1:],\n          tuple(other_transforms),\n      )\n    case _:\n      # No ExtractAliasedRef found, don't do anything.\n      return ref, ref_aval, transform_avals, transforms\n\n\ndef _commute_transform(\n    aval: jax_core.AbstractValue,\n    t1: state_types.Transform,\n    t2: state_types.Transform,\n) -> tuple[state_types.Transform, state_types.Transform]:\n  \"\"\"Commutes two transforms.\n","sourceCodeStart":1652,"sourceCodeEnd":1688,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L1652-L1688","documentation":"Terminal fallback in _extract_aliased_ref: the aliased ref's memory space is neither SMEM (smem()) nor TMEM (tmem()), so no aliasing strategy exists. The lowering switches on the ref's memory space and raises NotImplementedError for anything else, such as global/device memory or WMEM.","triggerScenarios":"Creating an aliased/viewed Ref that resolves to a memory space other than shared or tensor memory — e.g. aliasing a global-memory buffer or a warpgroup-register (wmem) ref inside a Pallas GPU kernel.","commonSituations":"Kernel code that aliases input/output refs directly (which live in global memory) instead of staging through SMEM; new memory spaces introduced by newer Mosaic dialects.","solutions":["Copy the global/wmem data into SMEM scratch first, then alias/bitcast the SMEM buffer","Remove view/aliasing from refs that live in unsupported memory spaces; convert dtypes on values instead","Check that the ref you're viewing was actually allocated by SMEM/TMEM helpers and not passed in from global memory"],"exampleFix":"# before\nscratch = in_ref.view(jnp.float32)  # in_ref is global memory -> error\n# after\nsmem = alloc_smem(in_ref.shape, jnp.uint8)\nsmem[...] = in_ref.astype(jnp.uint8)\nscratch = smem.view(jnp.float32)","handlingStrategy":"type-guard","validationCode":"space = getattr(buf, 'memory_space', None)\nassert space in ('smem', 'tmem'), f'cannot alias ref in memory space {space!r}'","typeGuard":"def aliasable_memory_space(buf) -> bool:\n    return getattr(buf, 'memory_space', None) in ('smem', 'tmem')","tryCatchPattern":null,"preventionTips":["Never call view() on kernel input/output refs (global memory)","Stage data through SMEM before reinterpretation","Keep a helper that only exposes aliasable buffers to view logic"],"tags":["jax","pallas","mosaic-gpu","memory-space","aliasing","not-implemented"],"backgroundTag":"unsupported-memory-space-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}