{"record":{"id":"54c255cb618f27c1","repo":"jax-ml/jax","slug":"with-memory-space-constraint-only-supports-hbm-vm","errorCode":null,"errorMessage":"with_memory_space_constraint only supports HBM, VMEM, SMEM, and HOST.","messagePattern":"with_memory_space_constraint only supports HBM, VMEM, SMEM, and HOST\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1079,"sourceCode":"  apply this to the arguments of a pallas_call and it will constrain them, but\n  other operations will not respect this constraint.\n\n  Args:\n    x: The array to constrain.\n    memory_space: The memory space to constrain to.\n\n  Returns:\n    The array ``x`` with the memory space constraint.\n  \"\"\"\n  if memory_space is pl_core.MemorySpace.ANY:\n    return x\n  if memory_space not in {\n      tpu_core.MemorySpace.HBM,\n      tpu_core.MemorySpace.VMEM,\n      tpu_core.MemorySpace.SMEM,\n      jax_core.MemorySpace.Host,\n  }:\n    raise NotImplementedError(\n        \"with_memory_space_constraint only supports HBM, VMEM, SMEM, and HOST.\"\n    )\n  return pl_core.with_memory_space_constraint_p.bind(\n      x, memory_space=memory_space)\n\n\ndef load(ref: Ref, *, mask: jax.Array | None = None) -> jax.Array:\n  \"\"\"Loads an array from the given ref.\n\n  If ``mask`` is not specified, this function has the same semantics as\n  ``ref[idx]`` in JAX.\n\n  Args:\n    ref: The ref to load from.\n    mask: An optional boolean mask specifying which indices to load.\n\n  Returns:\n    The loaded array.","sourceCodeStart":1061,"sourceCodeEnd":1097,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1061-L1097","documentation":"with_memory_space_constraint in JAX Mosaic Pallas annotates an array with a memory space constraint, but only HBM, VMEM, SMEM (TPU) and HOST are recognized. Passing any other memory space value raises NotImplementedError.","triggerScenarios":"Calling with_memory_space_constraint(x, memory_space) where memory_space is not one of tpu_core.MemorySpace.HBM/VMEM/SMEM or jax_core.MemorySpace.Host — e.g. a string like 'HBM', an ACC memory space, or an ASM space.","commonSituations":"Passing a string name instead of the enum; passing MemorySpace.ACCUMULATOR (use matmul_acc_lhs refs for accumulators); copy-pasting a GPU Pallas memory space enum into TPU code.","solutions":["Use tpu_core.MemorySpace.HBM/VMEM/SMEM or jax_core.MemorySpace.Host enum values, not strings","For accumulator memory use dedicated matmul primitives, not this API"],"exampleFix":"# before\nx = with_memory_space_constraint(x, 'HBM')\n# after\nfrom jax._src.interpreters import mlir\nx = with_memory_space_constraint(x, tpu_core.MemorySpace.HBM)","handlingStrategy":"type-guard","validationCode":"from jax._src.pallas.tpu import tpu_core as tpu_core_mod\n_allowed = {tpu_core_mod.MemorySpace.HBM, tpu_core_mod.MemorySpace.VMEM, tpu_core_mod.MemorySpace.SMEM}\nassert memory_space in _allowed or memory_space is jax_core.MemorySpace.Host","typeGuard":"def is_supported_memory_space(ms) -> bool:\n    return ms in {tpu_core.MemorySpace.HBM, tpu_core.MemorySpace.VMEM, tpu_core.MemorySpace.SMEM, jax_core.MemorySpace.Host}","tryCatchPattern":null,"preventionTips":["Always pass MemorySpace enums, never strings","Use ACC-space refs + matmul primitives for accumulators"],"tags":["jax","pallas","tpu","memory-space","unsupported-operation"],"backgroundTag":"unsupported-argument-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}