jax-ml/jax · error · ValueError

Found inconsistent memory spaces in multiref: {self.ref}

Error message

Found inconsistent memory spaces in multiref: {self.ref}

What it means

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.

Source

Thrown at jax/_src/state/types.py:405

  def swap(self, value, idx=()):
    from jax._src.state.primitives import ref_swap  # pyrefly: ignore[missing-import]
    return ref_swap(self, idx, value)

  def get(self, idx=()):
    from jax._src.state.primitives import ref_get  # pyrefly: ignore[missing-import]
    return ref_get(self, idx)

  @property
  def memory_space(self):
    def _mem_space(ref):
      if isinstance(ref, TransformedRef):
        return ref.memory_space
      return core.typeof(ref).memory_space if hasattr(ref, "aval") else ref.memory_space

    if self.multiref:
      ms, *rest = tuple(_mem_space(r) for r in self.ref)
      if not all(m == ms for m in rest):
        raise ValueError(
            f"Found inconsistent memory spaces in multiref: {self.ref}"
        )
      return ms

    return _mem_space(self.ref)

  def __getattr__(self, name):
    if self.multiref:
      return cast(MultiRefTransform, self.transforms[0]).getattr(name, self.ref)
    return getattr(self.ref, name)

  def __getitem__(self, slc):
    from jax._src.state.primitives import ref_get  # pyrefly: ignore[missing-import]
    return ref_get(self, slc)

  def __setitem__(self, slc, value):
    from jax._src.state.primitives import ref_set  # pyrefly: ignore[missing-import]
    return ref_set(self, slc, value)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate all refs in the group with the same memory_space
2. Split the group by memory space
3. Audit allocation sites for inconsistent memory_space kwargs

Example fix

# before
refs = (ref_a, ref_b)  # ref_b allocated with memory_space=Pinned
# after
ref_b = alloc(...)  # same memory_space as ref_a
refs = (ref_a, ref_b)

spaces = {getattr(core.typeof(r), "memory_space", None) for r in ref.ref}
assert len(spaces) == 1, f"mixed memory spaces: {spaces}"

When it happens

Trigger: 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.

Common situations: Mixing device-resident and host/pinned buffers into one select group; partial migration of allocations to a new memory space.

Related errors

• array ref with memory space only works inside of a `jit`.
• Memory space {self.memory_space} is not supported by mesh {s
• Cannot eagerly run with_memory_space_constraint.
• Unsupported core type: {core_type}
• Invalid memory space: {memory_space!r}

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