jax-ml/jax · error · NotImplementedError

with_memory_space_constraint only supports HBM, VMEM, SMEM,

Error message

with_memory_space_constraint only supports HBM, VMEM, SMEM, and HOST.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:1079

  apply this to the arguments of a pallas_call and it will constrain them, but
  other operations will not respect this constraint.

  Args:
    x: The array to constrain.
    memory_space: The memory space to constrain to.

  Returns:
    The array ``x`` with the memory space constraint.
  """
  if memory_space is pl_core.MemorySpace.ANY:
    return x
  if memory_space not in {
      tpu_core.MemorySpace.HBM,
      tpu_core.MemorySpace.VMEM,
      tpu_core.MemorySpace.SMEM,
      jax_core.MemorySpace.Host,
  }:
    raise NotImplementedError(
        "with_memory_space_constraint only supports HBM, VMEM, SMEM, and HOST."
    )
  return pl_core.with_memory_space_constraint_p.bind(
      x, memory_space=memory_space)


def load(ref: Ref, *, mask: jax.Array | None = None) -> jax.Array:
  """Loads an array from the given ref.

  If ``mask`` is not specified, this function has the same semantics as
  ``ref[idx]`` in JAX.

  Args:
    ref: The ref to load from.
    mask: An optional boolean mask specifying which indices to load.

  Returns:
    The loaded array.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use tpu_core.MemorySpace.HBM/VMEM/SMEM or jax_core.MemorySpace.Host enum values, not strings
2. For accumulator memory use dedicated matmul primitives, not this API

Example fix

# before
x = with_memory_space_constraint(x, 'HBM')
# after
from jax._src.interpreters import mlir
x = with_memory_space_constraint(x, tpu_core.MemorySpace.HBM)

from jax._src.pallas.tpu import tpu_core as tpu_core_mod
_allowed = {tpu_core_mod.MemorySpace.HBM, tpu_core_mod.MemorySpace.VMEM, tpu_core_mod.MemorySpace.SMEM}
assert memory_space in _allowed or memory_space is jax_core.MemorySpace.Host

def is_supported_memory_space(ms) -> bool:
    return ms in {tpu_core.MemorySpace.HBM, tpu_core.MemorySpace.VMEM, tpu_core.MemorySpace.SMEM, jax_core.MemorySpace.Host}

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Memory space {self.memory_space} is not supported by mesh {s
• Unsupported core type: {core_type}
• Invalid memory space: {memory_space!r}
• The Pallas TPU lowering currently supports in memory space A
• Loads are only allowed on VMEM and SMEM references.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/54c255cb618f27c1. Report an issue: GitHub.