{"record":{"id":"73d2eead08b360fc","repo":"jax-ml/jax","slug":"loads-and-stores-are-only-allowed-on-vmem-and-smem","errorCode":null,"errorMessage":"Loads and stores are only allowed on VMEM and SMEM references.{extra}","messagePattern":"Loads and stores are only allowed on VMEM and SMEM references\\.(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2525,"sourceCode":"  )\n  need_stride = not all((s is None or s == 1) for s in strides)\n\n  if is_smem_store:\n    if mask is not None:\n      raise ValueError(\"SMEM store does not support masks\")\n    if val_aval.shape:\n      raise ValueError(\"Can only store scalars to SMEM\")\n    result = memref.load(ref, starts)\n    result = _maybe_cast_load_to_bool(ctx, val_aval, result)\n    val = _maybe_cast_store_to_memref_type(ctx, val_aval, val)\n    memref.store(val, ref, starts)\n    return result\n\n  if not is_vmem_store:\n    extra = \"\"\n    if memory_space == \"#tpu.memory_space<any>\":\n      extra = \" ANY memory space can only be accessed using async_copy.\"\n    raise ValueError(\n        \"Loads and stores are only allowed on VMEM and SMEM references.\" + extra\n    )\n\n  # handling VMEM store below\n  if not val_aval.shape:\n    raise ValueError(\"Cannot store scalars to VMEM\")\n\n  mem_slice_shape = list(aval_out.shape)\n  for i, a in enumerate(idx.indices):\n    if not isinstance(a, primitives.Slice):\n      mem_slice_shape.insert(i, 1)\n  mem_slice_shape_iter = iter(mem_slice_shape)\n  mem_slice_shape = [\n      1 if b is pallas_core.squeezed else next(mem_slice_shape_iter)\n      for b in ref_block_shape\n  ]\n  mem_aval = aval_out.update(\n      shape=tuple(mem_slice_shape), sharding=jax_core.get_cur_mesh_sharding()","sourceCodeStart":2507,"sourceCodeEnd":2543,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2507-L2543","documentation":"Raised when a store targets a ref in a memory space other than VMEM or SMEM. Synchronous stores only work on vector and scalar memory; ANY-space buffers must be written via the async copy (DMA) path.","triggerScenarios":"pl.store to a ref whose memory_space is ANY (or otherwise not vmem/smem), e.g. trying to write final results directly from an ANY-space output without async_copy.","commonSituations":"Mixing async_copy pipelines with direct stores; writing to the ANY-space copy destination/source directly.","solutions":["Write into a VMEM scratch ref, then pltpu.async_copy(scratch, any_ref, ...) to move results out","Use pltpu.async_copy with the ANY-space buffer as the DMA source/destination only"],"exampleFix":"# before\npl.store(any_out_ref, val)\n# after\npl.store(vmem_scratch, val)\npltpu.async_copy(vmem_scratch, any_out_ref)\npltpu.async_copy_wait()","handlingStrategy":"fallback","validationCode":"def store_any_space(vmem_scratch, any_ref, val):\n    pl.store(vmem_scratch, val)\n    pltpu.async_copy(vmem_scratch, any_ref)\n    pltpu.async_copy_wait()","typeGuard":"def is_any_space(ref) -> bool:\n    return 'any>' in str(getattr(ref.aval, 'memory_space', ''))","tryCatchPattern":null,"preventionTips":["Write results to VMEM then async_copy out","Treat ANY space as DMA-only"],"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"}