{"record":{"id":"05cc4deaf5f10332","repo":"jax-ml/jax","slug":"pinned-array-ref-only-works-inside-of-a-jit","errorCode":null,"errorMessage":"pinned array ref only works inside of a `jit`.","messagePattern":"pinned array ref only works inside of a `jit`\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2969,"sourceCode":"def _ref_abstract_eval(init_aval, *, memory_space: Any, kind: Any, pin: bool):\n  from jax._src.state.types import AbstractRef  # pyrefly: ignore[missing-import]\n  # If no memory space is specified, use the memory space of the initial value\n  # but we make sure to reset it to Device because the Ref owns the memory space\n  if (memory_space is None\n      and isinstance(init_aval, ShapedArray)):\n    if init_aval.memory_space is not MemorySpace.Device:\n      memory_space = init_aval.memory_space\n    init_aval = init_aval.update(memory_space=MemorySpace.Device)\n  return (AbstractRef(init_aval, memory_space=memory_space, kind=kind),\n          {internal_mutable_array_effect})\n\n@ref_p.def_impl\ndef _ref_impl(init_val, *, memory_space: Any, kind: Any, pin: bool):\n  if memory_space is not None:\n    raise NotImplementedError(\n        \"array ref with memory space only works inside of a `jit`.\")\n  if pin:\n    raise NotImplementedError(\n        \"pinned array ref only works inside of a `jit`.\")\n  from jax._src.state.types import AbstractRef  # pyrefly: ignore[missing-import]\n  from jax._src.lax.lax import _array_copy  # pyrefly: ignore[missing-import]\n  aval = AbstractRef(typeof(init_val), kind=kind)\n  return Ref(aval, ArrayRefImpl(aval, _array_copy(init_val)))\n\n# TODO(mattjj,dougalm): merge with ref_p\ndef empty_ref(ty, memory_space=None, pin=False):\n  aval = shaped_abstractify(ty)\n  return empty_ref_p.bind(ty=aval, memory_space=memory_space, pin=pin)\nempty_ref_p = Primitive('empty_ref')\nempty_ref_p.ref_primitive = True\nempty_ref_p.is_effectful = lambda _: True\nempty_ref_p.ref_allocating = True\nempty_ref_p.is_high = lambda *, ty, memory_space, pin: ty.is_high\n\ndef _empty_ref_to_lojax(*, ty, memory_space, pin):\n  from jax._src.state.types import AbstractRef  # pyrefly: ignore[missing-import]","sourceCodeStart":2951,"sourceCodeEnd":2987,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2951-L2987","documentation":"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.","triggerScenarios":"Eagerly creating a pinned ref (pin=True) via jax state/ref primitives instead of inside a jitted function.","commonSituations":"Experimenting with experimental pinned/ref APIs outside jit; partial-trace debugging of stateful jitted code.","solutions":["Move the pinned ref creation inside jax.jit","Omit pin=True in eager code","Update JAX if eager support was added"],"exampleFix":"// before\nref = make_ref(init, pin=True)  # eager\n\n// after\nref = jax.jit(lambda v: make_ref(v, pin=True))(init)","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    ref = make_ref(v, pin=True)\nexcept NotImplementedError:\n    ref = jax.jit(lambda x: make_ref(x, pin=True))(v)","preventionTips":["Create pinned refs only within jitted functions"],"tags":["jax","refs","pinned-memory","jit","not-implemented"],"backgroundTag":"requires-jit-context","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}