jax-ml/jax · error · NotImplementedError

Unfolding dimensions is not supported when commuting an `Un

Error message

Unfolding dimensions is not supported when commuting an  `UntilingTransform` with a `ReshapeTransform`

What it means

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.

Source

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

    if not self.tiling:
      raise NotImplementedError(
          "Commuting a `UntilingTransform` with a `ReshapeTransform` is not "
          "supported when the tiling is empty"
      )
    untiled_aval = self.transform_type(aval)
    assert isinstance(untiled_aval, jax_core.ShapedArray)
    components = [[]]
    # We assume that we support only folds here for the moment. Therefore, we
    # can gather a number of consecutive dimensions such that their product
    # equals the dimension currently being processed in the reshaped shape.
    for d in untiled_aval.shape:
      reshaped_dim_size = transform.shape[len(components) - 1]
      components[-1].append(d)
      component_size = math.prod(components[-1])
      if component_size == reshaped_dim_size:
        components.append([])
      elif component_size > reshaped_dim_size:
        raise NotImplementedError(
            "Unfolding dimensions is not supported when commuting an "
            " `UntilingTransform` with a `ReshapeTransform`"
        )
    assert not components[-1]
    components.pop()
    assert len(components) == len(transform.shape)

    rev_tiling_to_process = list(self.tiling)[::-1]
    rev_shape_to_process = untiled_aval.shape[-len(self.tiling):][::-1]
    rev_new_tiling: list[int] = []
    rev_new_tiled_dims: list[int] = []
    for component in components[::-1]:
      # The construction above should guarantee that there is never an empty
      # component, which simplifies indexing below.
      assert component
      ndim = len(component)
      if len(rev_tiling_to_process) < ndim:
        raise NotImplementedError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the reshape fold-only: target dims should combine whole source dims, never split one
2. Insert an explicit untile (materialize) before the reshape and re-tile afterwards if needed
3. Align reshape boundaries with the tiling structure of the reference

Example fix

// before
x = x.reshape((4, 8))  # unfolds dim 32

// after
x = x.reshape((32, 1))  # fold-only, or materialize: x = jnp.asarray(x).reshape(4, 8)

import math
def is_fold_only(before, after):
    # each target dim must be a product of whole prefix source dims
    ...

try: ref.reshape(target)\nexcept NotImplementedError: ref = jnp.asarray(ref).reshape(target)

When it happens

Trigger: 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.

Common situations: Unflattening flattened tensors (view(-1) then view(4,8)) in TPU pallas kernels; reshapes that don't respect tiling boundaries.

Related errors

• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Folding tiled dimensions into untiled dimensions is not supp
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Reshape shape {shape} is not divisible by swizzle elements {
• packed cannot be specified if layout is specified.

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