{"record":{"id":"81ef8691e0b5b8c7","repo":"jax-ml/jax","slug":"found-inconsistent-memory-spaces-in-multiref-sel","errorCode":null,"errorMessage":"Found inconsistent memory spaces in multiref: {self.ref}","messagePattern":"Found inconsistent memory spaces in multiref: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/types.py","lineNumber":405,"sourceCode":"  def swap(self, value, idx=()):\n    from jax._src.state.primitives import ref_swap  # pyrefly: ignore[missing-import]\n    return ref_swap(self, idx, value)\n\n  def get(self, idx=()):\n    from jax._src.state.primitives import ref_get  # pyrefly: ignore[missing-import]\n    return ref_get(self, idx)\n\n  @property\n  def memory_space(self):\n    def _mem_space(ref):\n      if isinstance(ref, TransformedRef):\n        return ref.memory_space\n      return core.typeof(ref).memory_space if hasattr(ref, \"aval\") else ref.memory_space\n\n    if self.multiref:\n      ms, *rest = tuple(_mem_space(r) for r in self.ref)\n      if not all(m == ms for m in rest):\n        raise ValueError(\n            f\"Found inconsistent memory spaces in multiref: {self.ref}\"\n        )\n      return ms\n\n    return _mem_space(self.ref)\n\n  def __getattr__(self, name):\n    if self.multiref:\n      return cast(MultiRefTransform, self.transforms[0]).getattr(name, self.ref)\n    return getattr(self.ref, name)\n\n  def __getitem__(self, slc):\n    from jax._src.state.primitives import ref_get  # pyrefly: ignore[missing-import]\n    return ref_get(self, slc)\n\n  def __setitem__(self, slc, value):\n    from jax._src.state.primitives import ref_set  # pyrefly: ignore[missing-import]\n    return ref_set(self, slc, value)","sourceCodeStart":387,"sourceCodeEnd":423,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/types.py#L387-L423","documentation":"TransformedRef.memory_space raises ValueError when refs in a multiref live in different memory spaces (e.g. some in Device memory, some in pinned/host memory), since a single memory space can't be reported.","triggerScenarios":"Grouping refs allocated with different memory_space arguments into one multiref and then querying .memory_space or using it in operations that need a consistent space.","commonSituations":"Mixing device-resident and host/pinned buffers into one select group; partial migration of allocations to a new memory space.","solutions":["Allocate all refs in the group with the same memory_space","Split the group by memory space","Audit allocation sites for inconsistent memory_space kwargs"],"exampleFix":"# before\nrefs = (ref_a, ref_b)  # ref_b allocated with memory_space=Pinned\n# after\nref_b = alloc(...)  # same memory_space as ref_a\nrefs = (ref_a, ref_b)","handlingStrategy":"validation","validationCode":"spaces = {getattr(core.typeof(r), \"memory_space\", None) for r in ref.ref}\nassert len(spaces) == 1, f\"mixed memory spaces: {spaces}\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Standardize memory_space across grouped allocations","Add allocation factory helpers that fix memory_space"],"tags":["jax","memory-space","multiref"],"backgroundTag":"inconsistent-resource-config","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}