{"record":{"id":"c658c9b1af467e2e","repo":"jax-ml/jax","slug":"cannot-transform-type-x","errorCode":null,"errorMessage":"Cannot transform type: {x}","messagePattern":"Cannot transform type: (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":743,"sourceCode":"  \"\"\"\n  tiling: tuple[int, ...]\n\n  def transform_type(self, x):\n    match x:\n      case jax_core.ShapedArray():\n        shape = x.shape\n        if shape is None:\n          return x\n        leading_dims = shape[: -len(self.tiling) :]\n        tiled_dims = shape[-len(self.tiling) :]\n        assert all(d % t == 0 for d, t in zip(tiled_dims, self.tiling))\n        num_tiles = [d // t for d, t in zip(tiled_dims, self.tiling)]\n        new_shape = (*leading_dims, *num_tiles, *self.tiling)\n        return x.update(shape=new_shape)\n      case state_types.AbstractRef():\n        return x.update(inner_aval=self.transform_type(x.inner_aval))\n      case _:\n        raise TypeError(f\"Cannot transform type: {x}\")\n\n  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:\n    return UntilingTransform(self.tiling)\n\n@tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True)\nclass UntilingTransform(state_types.Transform):\n  tiling: tuple[int, ...] = jax.tree.static()\n\n  def transform_type(self, x):\n    match x:\n      case jax_core.ShapedArray():\n        shape = x.shape\n        if shape is None:\n          return x\n        assert shape[-len(self.tiling) :] == self.tiling, (shape, self.tiling)\n        shape = shape[: -len(self.tiling)]  # Drop tiling\n        new_shape = shape[: -len(self.tiling)] + tuple(","sourceCodeStart":725,"sourceCodeEnd":761,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L725-L761","documentation":"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}'.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Inspect the failing aval's type and ensure only arrays/refs reach the transform","Filter or map the pytree so non-array leaves bypass the tiling transform","If a custom aval must be supported, subclass/extend the transform's match statement upstream in your own fork"],"exampleFix":"# before\ntransformed = transform.transform_type(aval)  # TypeError on tokens\n\n# after\nfrom jax._src import state_types\nif isinstance(aval, (jax.core.ShapedArray, state_types.AbstractRef)):\n    transformed = transform.transform_type(aval)\nelse:\n    transformed = aval  # pass through untouched","handlingStrategy":"type-guard","validationCode":"from jax._src import state_types\nok = isinstance(aval, (jax.core.ShapedArray, state_types.AbstractRef))","typeGuard":"def is_transformable_aval(x) -> bool:\n    from jax._src import state_types\n    return isinstance(x, (jax.core.ShapedArray, state_types.AbstractRef))","tryCatchPattern":"null","preventionTips":["Keep tokens and non-array leaves out of block specs","Pre-validate pytree leaves before applying transforms"],"tags":["jax","pallas","abstract-values","transforms","type-dispatch"],"backgroundTag":"unsupported-type-dispatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}