{"record":{"id":"6e17520415fcfe73","repo":"jax-ml/jax","slug":"aval-memory-space-inner-aval-memory-space-does-n","errorCode":null,"errorMessage":"Aval memory space {inner_aval.memory_space} does not match the requested memory space {memory_space}.","messagePattern":"Aval memory space (.+?) does not match the requested memory space (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/types.py","lineNumber":471,"sourceCode":"  \"\"\"Abstract mutable array reference.\n\n  Refer to the `Ref guide`_ for more information.\n\n  .. _Ref guide: https://docs.jax.dev/en/latest/array_refs.html\n  \"\"\"\n  __slots__ = [\"inner_aval\", \"memory_space\", \"kind\"]\n\n  def __init__(self, inner_aval: core.AbstractValue, memory_space: Any = None,\n               kind: Any = None):\n    self.inner_aval = inner_aval\n    # TODO(sharadmv,mattjj,yashkatariya): merge memory spaces\n    if isinstance(inner_aval, core.ShapedArray):\n      if (\n          inner_aval.memory_space is not None\n          and inner_aval.memory_space != core.MemorySpace.Device\n          and inner_aval.memory_space != memory_space\n      ):\n        raise ValueError(\n            f\"Aval memory space {inner_aval.memory_space} does not match the \"\n            f\"requested memory space {memory_space}.\"\n        )\n      # Hide the inner_aval memory space.\n      self.inner_aval = inner_aval.update(memory_space=core.MemorySpace.Device)\n    self.memory_space = memory_space\n    self.kind = kind\n\n  @property\n  def is_high(self):\n    return self.inner_aval.is_high\n\n  def lo_ty(self):\n    return [self.update(inner_aval=x) for x in self.inner_aval.lo_ty()]\n\n  def lower_val(self, ref, /):\n    if not self.is_high:\n      return [ref]","sourceCodeStart":453,"sourceCodeEnd":489,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/types.py#L453-L489","documentation":"AbstractRef.__init__ (aval construction) raises ValueError when the inner aval already declares a non-default memory space that differs from the explicitly requested memory_space. The two must agree or the aval must use Device (default) space.","triggerScenarios":"Constructing a ref aval with memory_space=X while the wrapped ShapedArray carries memory_space=Y (both non-Device, X != Y).","commonSituations":"Wrapping an aval that was already placed in pinned memory and requesting device memory, or vice versa; double-application of memory-space logic after API changes.","solutions":["Pass the same memory_space that the inner aval already has","Reset the inner aval's memory_space to Device before re-wrapping","Trace which layer set the inner aval's memory_space and consolidate"],"exampleFix":"# before\nref = AbstractRef(inner_aval_with_pinned, memory_space=Device)\n# after\nref = AbstractRef(inner_aval.update(memory_space=core.MemorySpace.Device), memory_space=core.MemorySpace.Device)","handlingStrategy":"validation","validationCode":"ms = getattr(inner_aval, \"memory_space\", None)\nif ms is not None and ms != core.MemorySpace.Device and ms != memory_space:\n    inner_aval = inner_aval.update(memory_space=core.MemorySpace.Device)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never re-wrap avals with a conflicting memory_space","Centralize memory-space decisions in one layer"],"tags":["jax","memory-space","aval"],"backgroundTag":"config-conflict-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}