jax-ml/jax · error · TypeError

Cannot transform type: {x}

Error message

Cannot transform type: {x}

What it means

UntilingTransform.transform_type applies an untiling reshape to abstract values. It supports ShapedArray (rewriting shape into leading dims + tiled dims) and state_types.AbstractRef (recursing into inner_aval); any other AbstractValue type raises TypeError 'Cannot transform type: {x}'.

Source

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

  """
  tiling: tuple[int, ...]

  def transform_type(self, x):
    match x:
      case jax_core.ShapedArray():
        shape = x.shape
        if shape is None:
          return x
        leading_dims = shape[: -len(self.tiling) :]
        tiled_dims = shape[-len(self.tiling) :]
        assert all(d % t == 0 for d, t in zip(tiled_dims, self.tiling))
        num_tiles = [d // t for d, t in zip(tiled_dims, self.tiling)]
        new_shape = (*leading_dims, *num_tiles, *self.tiling)
        return x.update(shape=new_shape)
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case _:
        raise TypeError(f"Cannot transform type: {x}")

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    return UntilingTransform(self.tiling)

@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class UntilingTransform(state_types.Transform):
  tiling: tuple[int, ...] = jax.tree.static()

  def transform_type(self, x):
    match x:
      case jax_core.ShapedArray():
        shape = x.shape
        if shape is None:
          return x
        assert shape[-len(self.tiling) :] == self.tiling, (shape, self.tiling)
        shape = shape[: -len(self.tiling)]  # Drop tiling
        new_shape = shape[: -len(self.tiling)] + tuple(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect the failing aval's type and ensure only arrays/refs reach the transform
2. Filter or map the pytree so non-array leaves bypass the tiling transform
3. If a custom aval must be supported, subclass/extend the transform's match statement upstream in your own fork

Example fix

# before
transformed = transform.transform_type(aval)  # TypeError on tokens

# after
from jax._src import state_types
if isinstance(aval, (jax.core.ShapedArray, state_types.AbstractRef)):
    transformed = transform.transform_type(aval)
else:
    transformed = aval  # pass through untouched

from jax._src import state_types
ok = isinstance(aval, (jax.core.ShapedArray, state_types.AbstractRef))

def is_transformable_aval(x) -> bool:
    from jax._src import state_types
    return isinstance(x, (jax.core.ShapedArray, state_types.AbstractRef))

null

When it happens

Trigger: Calling transform_type on an aval that is neither a ShapedArray nor an AbstractRef — e.g. a token, DShapedArray, or custom abstract value appearing in a block mapping / ref aval during to_block_mapping or get_ref_aval.

Common situations: Extending Pallas pipelines with new aval types; passing pytrees containing tokens or non-array leaves through block spec transforms; JAX version changes introducing new abstract value kinds into state primitives.

Related errors

• Transpose cannot be moved before a tiling transform when it
• Unsupported index type: {type(idx)}
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unsupported transform: {type(transform)}

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