jax-ml/jax · error · TypeError

Unsupported type: {x}

Error message

Unsupported type: {x}

What it means

The transform_type of ExtractAliasedRef only supports AbstractRef (recursing into inner_aval) and ShapedArray inputs; any other abstract value type raises TypeError with the offending type name.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1164

  @classmethod
  def from_transformed_ref(
      cls,
      ref: pallas_core.TransformedRef,
      byte_offset: int,
      alias_group_idx: int,
      layout: tcgen05.TMEMLayout | None = None,
  ):
    return cls(dtypes.dtype(ref.dtype), ref.ref.shape, byte_offset, alias_group_idx, layout)

  def transform_type(self, x):
    match x:
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case jax_core.ShapedArray():
        return x.update(shape=self.shape, dtype=self.dtype)
      case _:
        raise TypeError(f"Unsupported type: {x}")


@dataclasses.dataclass(frozen=True)
class SwizzleTransform(state_types.Transform):
  swizzle: int

  def __post_init__(self):
    if self.swizzle not in {32, 64, 128}:
      raise ValueError(
          f"Swizzle {self.swizzle} is not supported. Only 32, 64 and 128 are"
          " accepted."
      )

  def transform_type(
      self, x: jax_core.AbstractValue
  ) -> jax_core.AbstractValue:
    match x:
      case jax_core.ShapedArray():

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure only array-like refs/blocks reach this path
2. Inspect the aval type (type(x)) in a debugger to find what non-array value is leaking in
3. Upgrade jax if the aval type is newly introduced and support was added later

assert isinstance(x, (jax_core.ShapedArray, state_types.AbstractRef))

def supported_aval(x): return isinstance(x, (jax_core.ShapedArray, state_types.AbstractRef))

When it happens

Trigger: Applying an extract-aliased-ref transform to an abstract value that is neither a ShapedArray nor an AbstractRef, e.g. a token or a differently-shaped aval, via get_ref_aval or block mapping conversion.

Common situations: Unusual avals flowing through pallas pipelines (tokens, custom avals); version mismatches where new aval kinds reach this transform.

Related errors

• ref must be a reference
• Can only store to references (got {x_ref}).
• Can only store scalars or vectors (got {value}).
• Can only store to references (got {x_smem}).
• MemoryRef type must be a ShapedArray, got {type(self.inner_a

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