{"record":{"id":"3187543f78833585","repo":"jax-ml/jax","slug":"unfolding-dimensions-is-not-supported-when-commuti","errorCode":null,"errorMessage":"Unfolding dimensions is not supported when commuting an  `UntilingTransform` with a `ReshapeTransform`","messagePattern":"Unfolding dimensions is not supported when commuting an  `UntilingTransform` with a `ReshapeTransform`","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":878,"sourceCode":"    if not self.tiling:\n      raise NotImplementedError(\n          \"Commuting a `UntilingTransform` with a `ReshapeTransform` is not \"\n          \"supported when the tiling is empty\"\n      )\n    untiled_aval = self.transform_type(aval)\n    assert isinstance(untiled_aval, jax_core.ShapedArray)\n    components = [[]]\n    # We assume that we support only folds here for the moment. Therefore, we\n    # can gather a number of consecutive dimensions such that their product\n    # equals the dimension currently being processed in the reshaped shape.\n    for d in untiled_aval.shape:\n      reshaped_dim_size = transform.shape[len(components) - 1]\n      components[-1].append(d)\n      component_size = math.prod(components[-1])\n      if component_size == reshaped_dim_size:\n        components.append([])\n      elif component_size > reshaped_dim_size:\n        raise NotImplementedError(\n            \"Unfolding dimensions is not supported when commuting an \"\n            \" `UntilingTransform` with a `ReshapeTransform`\"\n        )\n    assert not components[-1]\n    components.pop()\n    assert len(components) == len(transform.shape)\n\n    rev_tiling_to_process = list(self.tiling)[::-1]\n    rev_shape_to_process = untiled_aval.shape[-len(self.tiling):][::-1]\n    rev_new_tiling: list[int] = []\n    rev_new_tiled_dims: list[int] = []\n    for component in components[::-1]:\n      # The construction above should guarantee that there is never an empty\n      # component, which simplifies indexing below.\n      assert component\n      ndim = len(component)\n      if len(rev_tiling_to_process) < ndim:\n        raise NotImplementedError(","sourceCodeStart":860,"sourceCodeEnd":896,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L860-L896","documentation":"While grouping dimensions to commute an untile past a reshape, a group of dimensions reached a size larger than the target reshaped dimension. That means the reshape unfolds a dimension (splits one dimension into several), which the commutation logic does not support.","triggerScenarios":"Reshaping a tiled block such that one original dimension maps to multiple target dimensions, e.g. reshape((4, 8)) from shape (32,) where tiling doesn't align, inside a pallas mosaic kernel that keeps transforms on the ref.","commonSituations":"Unflattening flattened tensors (view(-1) then view(4,8)) in TPU pallas kernels; reshapes that don't respect tiling boundaries.","solutions":["Make the reshape fold-only: target dims should combine whole source dims, never split one","Insert an explicit untile (materialize) before the reshape and re-tile afterwards if needed","Align reshape boundaries with the tiling structure of the reference"],"exampleFix":"// before\nx = x.reshape((4, 8))  # unfolds dim 32\n\n// after\nx = x.reshape((32, 1))  # fold-only, or materialize: x = jnp.asarray(x).reshape(4, 8)","handlingStrategy":"fallback","validationCode":"import math\ndef is_fold_only(before, after):\n    # each target dim must be a product of whole prefix source dims\n    ...","typeGuard":null,"tryCatchPattern":"try: ref.reshape(target)\\nexcept NotImplementedError: ref = jnp.asarray(ref).reshape(target)","preventionTips":["Prefer fold-only reshapes on transformed refs","Materialize (untile) before aggressive reshapes"],"tags":["jax","pallas","mosaic-gpu","reshape","unfolding"],"backgroundTag":"unsupported-transform-commutation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}