jax-ml/jax · error · ValueError

Cannot commute `UntilingTransform` with `ReshapeTransform` w

Error message

Cannot commute `UntilingTransform` with `ReshapeTransform` when any of the dimensions being collapsed other than the minormost one has non-unit tiling. Attempted to reshape tiled slice of shape {before} into shape {after}

What it means

When collapsing several dimensions into one during untile/reshape commutation, every collapsed dimension except the minormost must have unit tiling. If a non-minor collapsed dim is tiled, the memory layout cannot be preserved and a ValueError is raised with the offending before/after shapes.

Source

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

      rev_tiling_slice = rev_tiling_to_process[:ndim]
      rev_shape_slice = rev_shape_to_process[:ndim]
      new_tiling_dim = math.prod(rev_tiling_slice)
      num_elems = math.prod(rev_shape_slice)
      assert num_elems % new_tiling_dim == 0
      new_tiled_dim = num_elems // new_tiling_dim
      # If any other dimension than the minormost one has non-unit tiling, then
      # we cannot commute the reshape and untile transforms.
      #
      # Note that we could also support the case where we are collapsing
      # trailing tiled dimensions where the tile size is the dimension size
      # (i.e. there is a single tile).
      if any(t != 1 for t in rev_tiling_slice[1:]):
        before = (
            *[s // t for s, t in zip(rev_shape_slice, rev_tiling_slice)],
            *rev_tiling_slice,
        )
        after = (new_tiled_dim, new_tiling_dim)
        raise ValueError(
            "Cannot commute `UntilingTransform` with `ReshapeTransform` when "
            "any of the dimensions being collapsed other than the minormost "
            "one has non-unit tiling. Attempted to reshape tiled slice of "
            f"shape {before} into shape {after}"
        )
      rev_new_tiling.append(new_tiling_dim)
      rev_new_tiled_dims.append(new_tiled_dim)
      rev_tiling_to_process = rev_tiling_to_process[ndim:]
      rev_shape_to_process = rev_shape_to_process[ndim:]
      if not rev_tiling_to_process:
        break
    assert not rev_tiling_to_process
    assert not rev_shape_to_process
    new_tiling = tuple(rev_new_tiling[::-1])
    new_tiled_dims = tuple(rev_new_tiled_dims[::-1])
    new_shape = (
        *transform.shape[:len(components) - len(rev_new_tiling)],
        *new_tiled_dims,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Adjust the BlockSpec tiling so non-minormost dims of the collapsed group have tiling 1
2. Reshape only within one tiled dimension rather than across tiled dimensions
3. Materialize the block (copy out of the transformed ref) before the aggressive reshape

Example fix

// before
# tiling (2, 2), reshape (2,2)->(4,)
block = block.reshape((4,))

// after
# use tiling (1, 2) so only the minor dim is tiled
block = block.reshape((4,))

# before reshape, check non-minor collapsed dims have unit tiling
assert all(t == 1 for t in tiling[:-1]), 'non-minor dims must be untiled'

try: reshape\nexcept ValueError as e: adjust BlockSpec tiling

When it happens

Trigger: Reshaping a block whose tiled slice has non-unit tiling on a non-minormost collapsed dimension, e.g. shape (2,2) tiled (2,2) reshaped to (4,) / tiling (4).

Common situations: Swizzled or tiled layouts (common with TMA block specs with tiling like (8, 128)) combined with a flatten/reshape across the major dimension.

Related errors

• Folding tiled dimensions into untiled dimensions is not supp
• packed cannot be specified if layout is specified.
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is

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