{"record":{"id":"97d91a6ca410f657","repo":"jax-ml/jax","slug":"loads-are-only-allowed-on-vmem-and-smem-references","errorCode":null,"errorMessage":"Loads are only allowed on VMEM and SMEM references.","messagePattern":"Loads are only allowed on VMEM and SMEM references\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2343,"sourceCode":"    physical_out_shape = aval_out.shape\n  if not is_smem_load and not ref_block_shape:\n    raise NotImplementedError(\n        \"Indexing into a ()-shaped Ref not yet supported on TPU.\")\n  starts, sizes, strides, _, _ = _indexer_to_start_size_stride(\n      idx,\n      ref_block_shape,\n      cast_to_index=True,\n  )\n  need_stride = not all((s is None or s == 1) for s in strides)\n  if is_smem_load:\n    if ctx.avals_out[0].shape:\n      raise ValueError(\"Can only load scalars from SMEM\")\n    return _maybe_cast_load_to_bool(ctx, aval_out, memref.load(ref, starts))\n  elif str(ref_type.memory_space) != \"#tpu.memory_space<vmem>\":\n    extra = \"\"\n    if str(ref_type.memory_space) == \"#tpu.memory_space<any>\":\n      extra = \" ANY memory space can only be accessed using async_copy.\"\n    raise ValueError(\n        \"Loads are only allowed on VMEM and SMEM references.\" + extra\n    )\n  load_aval = jax_core.ShapedArray(sizes, dtype=physical_out_dtype)\n  if need_stride:\n    load_val = tpu.strided_load(\n        ctx.aval_to_ir_type(load_aval, is_kernel_boundary=True),\n        ref,\n        starts,\n        strides,\n    )\n  else:\n    load_val = vector.load(\n        ctx.aval_to_ir_type(load_aval, is_kernel_boundary=True),\n        ref,\n        starts,\n    )\n  if load_aval != aval_out:\n    if physical_out_shape:","sourceCodeStart":2325,"sourceCodeEnd":2361,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2325-L2361","documentation":"Raised when a Pallas load targets a ref in a memory space other than VMEM or SMEM. Only vector memory (VMEM) and scalar memory (SMEM) support synchronous loads; other spaces (notably ANY) require the asynchronous DMA copy path (async_copy).","triggerScenarios":"pl.load on a ref in e.g. '#tpu.memory_space<any>' or HBM-typed space instead of VMEM; using a BlockSpec that left memory space as ANY and reading it directly.","commonSituations":"Using pltpu.async_copy for DMA but also trying to pl.load the same ANY-space ref; forgetting to copy from ANY space into VMEM before compute.","solutions":["Use async_copy to move data from ANY space into a VMEM scratch ref, then load from that ref","Or specify memory_space=VMEM explicitly in the BlockSpec if synchronous load is intended"],"exampleFix":"# before\nx = pl.load(any_ref)  # any_ref has MemorySpace.ANY\n# after\nbuf = pltpu.make_partitioned(...)  # VMEM scratch\npltpu.async_copy(any_ref, buf, ...)\npltpu.async_copy_wait(...)\nx = pl.load(buf)","handlingStrategy":"fallback","validationCode":"def load_any_space(ref, vmem_scratch):\n    if str(getattr(ref.aval, 'memory_space', 'vmem')).endswith('any>'):\n        pltpu.async_copy(ref, vmem_scratch); pltpu.async_copy_wait()\n        return pl.load(vmem_scratch)\n    return pl.load(ref)","typeGuard":"def is_any_space(ref) -> bool:\n    return 'any>' in str(getattr(ref.aval, 'memory_space', ''))","tryCatchPattern":null,"preventionTips":["Never pl.load ANY-space refs; always async_copy into VMEM first","Structure kernels as DMA stage -> compute stage"],"tags":["jax","pallas","tpu","memory-space","async-copy"],"backgroundTag":"pallas-memory-space-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}