{"record":{"id":"af026d55ee69327c","repo":"jax-ml/jax","slug":"memoryref-type-must-be-a-shapedarray-got-type-se","errorCode":null,"errorMessage":"MemoryRef type must be a ShapedArray, got {type(self.inner_aval)}","messagePattern":"MemoryRef type must be a ShapedArray, got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":249,"sourceCode":"      pipeline expects the pre-populated buffer to be passed in via allocations\n      and will skip internal allocation.\n  \"\"\"\n  buffer_count: int\n  use_lookahead: bool = False\n  revisit: RevisitMode | None = None\n  prefetched_count: int = 0\n\n\n@dataclasses.dataclass(frozen=True)\nclass MemoryRef:\n  \"\"\"Like jax.ShapeDtypeStruct but with memory spaces.\"\"\"\n  inner_aval: jax_core.AbstractValue\n  # TODO(b/368122763): Unify memory space types across backends\n  memory_space: Any\n\n  def get_array_aval(self) -> jax_core.ShapedArray:\n    if not isinstance(self.inner_aval, jax_core.ShapedArray):\n      raise ValueError(\n          f\"MemoryRef type must be a ShapedArray, got {type(self.inner_aval)}\")\n    dtype = self.inner_aval.dtype\n    if not isinstance(dtype, (jnp.dtype, dtypes.ExtendedDType)):\n      dtype = jnp.dtype(dtype)\n    return self.inner_aval.update(dtype=dtype, memory_space=self.memory_space)\n\n  def get_ref_aval(self) -> TransformedRef | state.AbstractRef:\n    return state.AbstractRef(self.inner_aval, self.memory_space)\n\n  @property\n  def dtype(self):\n    return self.inner_aval.dtype  # pyrefly: ignore[missing-attribute]\n\n  @property\n  def shape(self):\n    return self.inner_aval.shape  # pyrefly: ignore[missing-attribute]\n\n  def __lt__(self, other):","sourceCodeStart":231,"sourceCodeEnd":267,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L231-L267","documentation":"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.","triggerScenarios":"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.","commonSituations":"Writing custom Pallas transformations or intercepting pallas_call outputs; JAX version upgrades that changed ref/aval internals.","solutions":["Inspect what aval is being wrapped (print type(self.inner_aval) via a debugger) and where it originates","Ensure refs passed to/from the kernel refer to array-valued specs, not tokens or sentinel avals","Update to a matching JAX version if this arises from internal plumbing rather than your code"],"exampleFix":null,"handlingStrategy":"try-catch","validationCode":null,"typeGuard":"def is_shaped_array_aval(aval):\n    from jax._src.core import ShapedArray\n    return isinstance(aval, ShapedArray)","tryCatchPattern":"wrap MemoryRef construction in try/except ValueError and assert the source aval is ShapedArray","preventionTips":["Only wrap array-valued specs in MemoryRef","Pin JAX version when using pallas internals"],"tags":["jax","pallas","aval","type-validation"],"backgroundTag":"invalid-type-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}