jax-ml/jax · error · ValueError

Aval memory space {inner_aval.memory_space} does not match t

Error message

Aval memory space {inner_aval.memory_space} does not match the requested memory space {memory_space}.

What it means

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.

Source

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

  """Abstract mutable array reference.

  Refer to the `Ref guide`_ for more information.

  .. _Ref guide: https://docs.jax.dev/en/latest/array_refs.html
  """
  __slots__ = ["inner_aval", "memory_space", "kind"]

  def __init__(self, inner_aval: core.AbstractValue, memory_space: Any = None,
               kind: Any = None):
    self.inner_aval = inner_aval
    # TODO(sharadmv,mattjj,yashkatariya): merge memory spaces
    if isinstance(inner_aval, core.ShapedArray):
      if (
          inner_aval.memory_space is not None
          and inner_aval.memory_space != core.MemorySpace.Device
          and inner_aval.memory_space != memory_space
      ):
        raise ValueError(
            f"Aval memory space {inner_aval.memory_space} does not match the "
            f"requested memory space {memory_space}."
        )
      # Hide the inner_aval memory space.
      self.inner_aval = inner_aval.update(memory_space=core.MemorySpace.Device)
    self.memory_space = memory_space
    self.kind = kind

  @property
  def is_high(self):
    return self.inner_aval.is_high

  def lo_ty(self):
    return [self.update(inner_aval=x) for x in self.inner_aval.lo_ty()]

  def lower_val(self, ref, /):
    if not self.is_high:
      return [ref]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the same memory_space that the inner aval already has
2. Reset the inner aval's memory_space to Device before re-wrapping
3. Trace which layer set the inner aval's memory_space and consolidate

Example fix

# before
ref = AbstractRef(inner_aval_with_pinned, memory_space=Device)
# after
ref = AbstractRef(inner_aval.update(memory_space=core.MemorySpace.Device), memory_space=core.MemorySpace.Device)

ms = getattr(inner_aval, "memory_space", None)
if ms is not None and ms != core.MemorySpace.Device and ms != memory_space:
    inner_aval = inner_aval.update(memory_space=core.MemorySpace.Device)

When it happens

Trigger: Constructing a ref aval with memory_space=X while the wrapped ShapedArray carries memory_space=Y (both non-Device, X != Y).

Common situations: 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.

Related errors

• linearized function called on tangent values inconsistent wi
• array ref with memory space only works inside of a `jit`.
• {self!r} has no `sharding`.
• MemoryRef type must be a ShapedArray, got {type(self.inner_a
• Memory space {self.memory_space} is not supported by mesh {s

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