{"record":{"id":"ff9b2a6ac32a775f","repo":"jax-ml/jax","slug":"gather-only-supports-loading-from-vmem-got-ref-a","errorCode":null,"errorMessage":"Gather only supports loading from VMEM, got {ref_aval.memory_space}","messagePattern":"Gather only supports loading from VMEM, got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_primitives.py","lineNumber":276,"sourceCode":"      raise ValueError(\n          f\"{mask.shape=} does not match the expected shape {out_aval.shape}\"\n      )\n    if mask.dtype != jnp.bool:\n      raise TypeError(f\"Mask must be a boolean array, got {mask.dtype}\")\n  return out_aval, {state_types.ReadEffect(0)}\n\n\n@sc_lowering.register_lowering_rule(gather_p)\ndef _gather_lowering_rule(\n    ctx: sc_lowering.LoweringRuleContext, *flat_args, tree\n):\n  ref, transforms, indices, mask = tree.unflatten(flat_args)\n  ref_aval, *_ = tree.unflatten(ctx.avals_in)\n  if ref_aval.memory_space not in (\n      tpu_core.MemorySpace.VMEM,\n      pallas_core.MemorySpace.DEFAULT,\n  ):\n    raise ValueError(\n        f\"Gather only supports loading from VMEM, got {ref_aval.memory_space}\"\n    )\n  if transforms:\n    ref_block_shape, *_ = ctx.block_shapes\n    ref, _ = tc_lowering._transform_ref(\n        ref, ref_aval, ref_block_shape, transforms\n    )\n  [out_aval] = ctx.avals_out\n  vec_type = ir.VectorType.get(\n      out_aval.shape, sc_lowering._dtype_to_ir_type(ref_aval.dtype)\n  )\n  return tpu.vector_load_idx(vec_type, ref, indices, mask=mask)\n\n\ndef load_gather(\n    ref: Ref, indices: Sequence[jax.Array], *, mask: jax.Array | None = None\n) -> jax.Array:\n  \"\"\"Gathers an array from a ref.","sourceCodeStart":258,"sourceCodeEnd":294,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_primitives.py#L258-L294","documentation":"Raised by the SparseCore lowering rule for pallas_load/gather when the reference being loaded from is not in VMEM (Vector Memory). Mosaic's SparseCore gather only supports reading out of VMEM or the DEFAULT memory space; any other memory space (e.g. SMEM/CMEM) is rejected at lowering time.","triggerScenarios":"Calling a Pallas kernel using load/gather on a tpu_core.BlockLayout ref whose memory_space is not tpu_core.MemorySpace.VMEM or pallas_core.MemorySpace.DEFAULT, e.g. a ref annotated with a CoreMemorySpace pointing at a non-VMEM space.","commonSituations":"Writing a SparseCore (SC) kernel and declaring intermediate scratch buffers or refs in the wrong memory space; porting a TensorCore kernel to SparseCore where SMEM was used; upgrading JAX where memory-space checking became stricter.","solutions":["Declare/annotate the gathered ref with memory space VMEM (tpu_core.MemorySpace.VMEM) or leave it DEFAULT","Check where the ref is created (kernel signature or scratch allocation) and fix its memory space","Avoid gather on non-VMEM buffers; copy to a VMEM buffer first inside the kernel"],"exampleFix":"// before\nref = pallas_core.new_memory_scope(...)  # memory_space=SMEM\nval = sc_primitives.load(ref, ...)  # gather path\n\n// after\n# allocate/cast the ref in VMEM\nval = kernel_call.load_from_vmem(ref, ...)","handlingStrategy":"validation","validationCode":"ms = ref.aval.memory_space if hasattr(ref, 'aval') else ref.memory_space\nassert ms in (tpu_core.MemorySpace.VMEM, pallas_core.MemorySpace.DEFAULT), ms","typeGuard":"def is_gatherable_ref(ref) -> bool:\n    ms = getattr(getattr(ref, 'aval', ref), 'memory_space', None)\n    return ms in (tpu_core.MemorySpace.VMEM, pallas_core.MemorySpace.DEFAULT)","tryCatchPattern":"try:\n    v = load_scatterlike(ref, indices)\nexcept ValueError as e:\n    if 'only supports loading from VMEM' in str(e):\n        raise RuntimeError(f'ref in unsupported memory space; move to VMEM: {e}')\n    raise","preventionTips":["Declare all gathered refs in VMEM in kernel signatures","Write a helper that allocates scratch only in VMEM for SC kernels","Add unit asserts on memory spaces when building kernels for multiple core types"],"tags":["jax","pallas","tpu","sparsecore","memory-space","gather"],"backgroundTag":"jax-pallas-memory-space-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}