{"record":{"id":"08daf52498690f22","repo":"jax-ml/jax","slug":"get-does-not-support-loading-from-ref-memory-spac","errorCode":null,"errorMessage":"Get does not support loading from {ref_memory_space!r}. Copy the data to a core-local memory space, e.g. VMEM, via `pltpu.async_copy`.","messagePattern":"Get does not support loading from (.+?)\\. Copy the data to a core-local memory space, e\\.g\\. VMEM, via `pltpu\\.async_copy`\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_lowering.py","lineNumber":90,"sourceCode":"  return _load_lowering_rule(ctx, ref, None, *flat_transforms, tree=tree)\n\n\ndef _load_lowering_rule(\n    ctx: LoweringRuleContext, ref, mask, *flat_transforms, tree\n):\n  ref_aval, *_flat_index_avals = ctx.avals_in\n  assert isinstance(ref_aval, state.AbstractRef)\n  [out_aval] = ctx.avals_out\n  assert isinstance(out_aval, jax_core.ShapedArray)\n\n  ref_memory_space = tpu_core.memory_space_to_tpu_memory_space(\n      ref_aval.memory_space, ctx.lowering_context.kernel_type\n  )\n  if (\n      ref_memory_space is MemorySpace.HBM\n      or ref_memory_space is MemorySpace.VMEM_SHARED\n  ):\n    raise NotImplementedError(\n        f\"Get does not support loading from {ref_memory_space!r}.\"\n        \" Copy the data to a core-local memory space, e.g. VMEM,\"\n        \" via `pltpu.async_copy`.\"\n    )\n\n  transforms = list(tree_util.tree_unflatten(tree, flat_transforms))\n  if not transforms or not isinstance(transforms[-1], indexing.NDIndexer):\n    tref_aval = state.transform_type(transforms, ref_aval)\n    assert isinstance(tref_aval, state.AbstractRef)\n    transforms.append(indexing.NDIndexer.make_trivial_indexer(tref_aval.shape))\n  *prev_transforms, indexer = transforms\n  ref_block_shape, *_ = ctx.block_shapes\n  ref, ref_block_shape = _transform_ref(\n      ref, ref_aval, ref_block_shape, prev_transforms\n  )\n  starts, sizes, strides, squeeze_dims, _ = tc_lowering._indexer_to_start_size_stride(\n      indexer, ref_block_shape, cast_to_index=True\n  )","sourceCodeStart":72,"sourceCodeEnd":108,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_lowering.py#L72-L108","documentation":"On SparseCore, the Get (ref read) lowering only reads core-local memory spaces. Loading directly from HBM or VMEM_SHARED refs is not implemented; data must first be moved with a DMA copy (pltpu.async_copy) into VMEM/SMEM.","triggerScenarios":"Reading an HSM/VMEM_SHARED ref with Python getitem (ref[...]) inside a SparseCore Pallas kernel, i.e. calling pallas_get on an HBM-backed ref.","commonSituations":"Porting a TPU TensorCore kernel that directly indexes HBM refs; forgetting the async_copy pipeline stage in SC kernels.","solutions":["Add an async_copy (DMA) from HBM into a VMEM scratch buffer and read from that","Restructure the kernel so reads come from SMEM/VMEM refs only"],"exampleFix":"# before\nv = hbm_ref[i, :]  # direct read from HBM in SC kernel\n# after\npltpu.async_copy(hbm_ref[i, :], vmem_buffer)  # in copy pipeline stage\nv = vmem_buffer[:]","handlingStrategy":"fallback","validationCode":"from jax._src.pallas.mosaic.tpu_core import MemorySpace\nif ref_aval.memory_space in (MemorySpace.HBM, MemorySpace.VMEM_SHARED):\n    raise SystemExit('stage via pltpu.async_copy first')","typeGuard":"null","tryCatchPattern":"null","preventionTips":["Structure SC kernels as copy-stage/compute-stage with VMEM buffers","Never index HBM refs directly in SC kernels"],"tags":["jax","pallas","tpu","sparsecore","memory-space"],"backgroundTag":"unsupported-memory-space-access","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}