jax-ml/jax · error · NotImplementedError

pinned array ref only works inside of a `jit`.

Error message

pinned array ref only works inside of a `jit`.

What it means

ref_p's eager implementation also rejects pinned array refs (pin=True): pinning is a compilation-time feature and only works inside jit. Eager calls with pin=True raise NotImplementedError.

Source

Thrown at jax/_src/core.py:2969

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

def _empty_ref_to_lojax(*, ty, memory_space, pin):
  from jax._src.state.types import AbstractRef  # pyrefly: ignore[missing-import]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move the pinned ref creation inside jax.jit
2. Omit pin=True in eager code
3. Update JAX if eager support was added

Example fix

// before
ref = make_ref(init, pin=True)  # eager

// after
ref = jax.jit(lambda v: make_ref(v, pin=True))(init)

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

When it happens

Trigger: Eagerly creating a pinned ref (pin=True) via jax state/ref primitives instead of inside a jitted function.

Common situations: Experimenting with experimental pinned/ref APIs outside jit; partial-trace debugging of stateful jitted code.

Related errors

• array ref with memory space only works inside of a `jit`.
• In order to best JIT compile `mode`, we cannot know whether
• In order to best JIT compile `rankdata`, we cannot know whet
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which

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