jax-ml/jax · error · ValueError

MemoryRef type must be a ShapedArray, got {type(self.inner_a

Error message

MemoryRef type must be a ShapedArray, got {type(self.inner_aval)}

What it means

A Pallas MemoryRef wraps an inner abstract value (aval) that must be a jax_core.ShapedArray; get_array_aval converts it to a concrete array aval. If a ref was created over a non-array aval (e.g. a token or a bare abstract value), the conversion raises ValueError because no shape/dtype semantics exist.

Source

Thrown at jax/_src/pallas/core.py:249

      pipeline expects the pre-populated buffer to be passed in via allocations
      and will skip internal allocation.
  """
  buffer_count: int
  use_lookahead: bool = False
  revisit: RevisitMode | None = None
  prefetched_count: int = 0


@dataclasses.dataclass(frozen=True)
class MemoryRef:
  """Like jax.ShapeDtypeStruct but with memory spaces."""
  inner_aval: jax_core.AbstractValue
  # TODO(b/368122763): Unify memory space types across backends
  memory_space: Any

  def get_array_aval(self) -> jax_core.ShapedArray:
    if not isinstance(self.inner_aval, jax_core.ShapedArray):
      raise ValueError(
          f"MemoryRef type must be a ShapedArray, got {type(self.inner_aval)}")
    dtype = self.inner_aval.dtype
    if not isinstance(dtype, (jnp.dtype, dtypes.ExtendedDType)):
      dtype = jnp.dtype(dtype)
    return self.inner_aval.update(dtype=dtype, memory_space=self.memory_space)

  def get_ref_aval(self) -> TransformedRef | state.AbstractRef:
    return state.AbstractRef(self.inner_aval, self.memory_space)

  @property
  def dtype(self):
    return self.inner_aval.dtype  # pyrefly: ignore[missing-attribute]

  @property
  def shape(self):
    return self.inner_aval.shape  # pyrefly: ignore[missing-attribute]

  def __lt__(self, other):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect what aval is being wrapped (print type(self.inner_aval) via a debugger) and where it originates
2. Ensure refs passed to/from the kernel refer to array-valued specs, not tokens or sentinel avals
3. Update to a matching JAX version if this arises from internal plumbing rather than your code

def is_shaped_array_aval(aval):
    from jax._src.core import ShapedArray
    return isinstance(aval, ShapedArray)

wrap MemoryRef construction in try/except ValueError and assert the source aval is ShapedArray

When it happens

Trigger: Constructing MemoryRef with an inner_aval that is not ShapedArray — usually indirectly via BlockSpec/indexing plumbing or custom pallas_call wrappers that pass wrong avals into _convert_out_shape_to_aval.

Common situations: Writing custom Pallas transformations or intercepting pallas_call outputs; JAX version upgrades that changed ref/aval internals.

Related errors

• Unsupported block dimension type: {type(dim)}. Allowed types
• Array type must have a `shape` attribute, but got {type(arra
• Unsupported block dimension type: {type(bd)}
• DMA source/destination/semaphore arguments must be Refs.
• DMA source semaphore must be a Ref.

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