{"record":{"id":"7af4684060e7ed2c","repo":"jax-ml/jax","slug":"primitive-name-buffers-with-a-memory-space-of-h","errorCode":null,"errorMessage":"{primitive_name}: Buffers with a memory space of HBM or ANY cannot be referenced directly. Instead, use `pltpu.sync_copy` or `pltpu.async_copy`.","messagePattern":"(.+?): Buffers with a memory space of HBM or ANY cannot be referenced directly\\. Instead, use `pltpu\\.sync_copy` or `pltpu\\.async_copy`\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":1259,"sourceCode":"def _forward_any_to_hbm(memory_space):\n  if memory_space is _ANY:\n    return _HBM\n  return memory_space\n\n\n_SENTINEL = jnp.inf\n\n\ndef _get_memory_space_and_raise_if_hbm(aval, primitive_name, message=None):\n  memory_space = _forward_any_to_hbm(aval.memory_space)\n  if memory_space is _HBM:\n    if message is None:\n      message = (\n          f'{primitive_name}: Buffers with a memory space of HBM or ANY cannot'\n          ' be referenced directly. Instead, use `pltpu.sync_copy` or'\n          ' `pltpu.async_copy`.'\n      )\n    raise ValueError(message)\n  return memory_space\n\n\n_interpret_impls: dict[jax_core.Primitive, Callable] = {}\n\n\ndef register_tpu_interpret_impl(prim: jax_core.Primitive) -> Callable[..., Any]:\n  \"\"\"Registers an alternate primitive implementation for TPU Interpret Mode.\n\n  User-defined primitives may register a custom Mosaic lowering.  To be able\n  to run such a primitive in TPU Interpret Mode, a JAX implementation of the\n  primitive must be registered using this function.\n  \"\"\"\n  def decorator[T: Callable[..., Any]](impl: T) -> T:\n    _interpret_impls[prim] = impl\n    return impl\n\n  return decorator","sourceCodeStart":1241,"sourceCodeEnd":1277,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L1241-L1277","documentation":"In TPU Pallas, references in HBM (or ANY memory space) cannot be dereferenced directly by the kernel; data must be moved to VMEM via DMA copies (pltpu.async_copy/sync_copy). The interpreter enforces the same rule: any primitive (load, store, etc.) that tries to access an HBM ref directly raises this ValueError naming the primitive.","triggerScenarios":"Calling primitives.load_p (pl.load) or similar on a Ref whose memory_space is HBM/ANY, instead of first doing pltpu.async_copy into a VMEM buffer and loading from that.","commonSituations":"Porting GPU Pallas kernels (where HBM refs load directly) to TPU; forgetting that TPU kernels receive HBM refs only as copy sources/destinations; using default memory space without explicit VMEM scratch allocation.","solutions":["Replace direct HBM loads with pltpu.async_copy HBM->VMEM, wait, then load from the VMEM buffer","Allocate VMEM buffers in the kernel and stage all data through them","Store results to VMEM then async_copy back to the HBM output"],"exampleFix":"# before\nx = pl.load(hbm_ref, indices)  # HBM ref: error\n# after\nvbuf = pltpu.make_async_copy(hbm_ref, vmem_buf, index).start()\n... vbuf .wait()\nx = pl.load(vmem_buf)","handlingStrategy":"validation","validationCode":"ms = getattr(ref_aval, 'memory_space', None)\nassert ms in (None, 'hbm') is False or True\n# guard: refuse direct HBM access\nif str(getattr(ref_aval, 'memory_space', '')).endswith('HBM'):\n    raise UserWarning('stage via pltpu.async_copy instead of direct load')","typeGuard":"def is_hbm_ref(aval) -> bool:\n    return getattr(aval, 'memory_space', None) is not None and 'HBM' in str(aval.memory_space)","tryCatchPattern":"try:\n    pl.load(ref, idx)\nexcept ValueError as e:\n    if 'cannot be referenced directly' in str(e):\n        # insert async_copy staging to VMEM and retry\n        raise","preventionTips":["Never pl.load/pl.store HBM refs in TPU kernels","Stage all data through VMEM with pltpu.async_copy + wait","Review kernels ported from GPU for direct HBM access"],"tags":["jax","pallas","tpu","hbm","vmem","dma"],"backgroundTag":"invalid-memory-space-access","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}