jax-ml/jax · error · NotImplementedError

array ref with memory space only works inside of a `jit`.

Error message

array ref with memory space only works inside of a `jit`.

What it means

ref_p's eager implementation cannot create an array ref with a memory_space — memory spaces are only handled by the jit-compiled path (where buffers can be allocated in specific memories). Calling the primitive eagerly with memory_space set raises NotImplementedError.

Source

Thrown at jax/_src/core.py:2966

ref_p.to_lojax = _ref_to_lojax

@ref_p.def_effectful_abstract_eval
def _ref_abstract_eval(init_aval, *, memory_space: Any, kind: Any, pin: bool):
  from jax._src.state.types import AbstractRef  # pyrefly: ignore[missing-import]
  # If no memory space is specified, use the memory space of the initial value
  # but we make sure to reset it to Device because the Ref owns the memory space
  if (memory_space is None
      and isinstance(init_aval, ShapedArray)):
    if init_aval.memory_space is not MemorySpace.Device:
      memory_space = init_aval.memory_space
    init_aval = init_aval.update(memory_space=MemorySpace.Device)
  return (AbstractRef(init_aval, memory_space=memory_space, kind=kind),
          {internal_mutable_array_effect})

@ref_p.def_impl
def _ref_impl(init_val, *, memory_space: Any, kind: Any, pin: bool):
  if memory_space is not None:
    raise NotImplementedError(
        "array ref with memory space only works inside of a `jit`.")
  if pin:
    raise NotImplementedError(
        "pinned array ref only works inside of a `jit`.")
  from jax._src.state.types import AbstractRef  # pyrefly: ignore[missing-import]
  from jax._src.lax.lax import _array_copy  # pyrefly: ignore[missing-import]
  aval = AbstractRef(typeof(init_val), kind=kind)
  return Ref(aval, ArrayRefImpl(aval, _array_copy(init_val)))

# TODO(mattjj,dougalm): merge with ref_p
def empty_ref(ty, memory_space=None, pin=False):
  aval = shaped_abstractify(ty)
  return empty_ref_p.bind(ty=aval, memory_space=memory_space, pin=pin)
empty_ref_p = Primitive('empty_ref')
empty_ref_p.ref_primitive = True
empty_ref_p.is_effectful = lambda _: True
empty_ref_p.ref_allocating = True
empty_ref_p.is_high = lambda *, ty, memory_space, pin: ty.is_high

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the reference-creating code in jax.jit so the compiled path handles memory_space
2. Drop the memory_space argument in eager paths
3. Wait for/use a JAX version implementing eager memory-space refs

Example fix

// before
ref = make_ref(init, memory_space='pinned')  # eager

// after
make_ref_jit = jax.jit(lambda v: make_ref(v, memory_space='pinned'))
ref = make_ref_jit(init)

try:
    ref = make_ref(v, memory_space=ms)
except NotImplementedError:
    ref = jax.jit(lambda x: make_ref(x, memory_space=ms))(v)

When it happens

Trigger: Invoking jax._src.state reference APIs (ref_p / experimental state primitives) with a memory_space argument outside of jit, e.g. via a stripped trace or direct primitive call in eager code.

Common situations: Using experimental jax.state/ref APIs directly; debugging/decomposing jitted stateful functions and re-executing pieces eagerly.

Related errors

• pinned array ref only works inside of a `jit`.
• Unsupported memory space: {x.memory_space}
• Only SMEM and TMEM refs are supported.
• Unsupported memory space.
• In order to best JIT compile `mode`, we cannot know whether

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