jax-ml/jax · error · NotImplementedError

Folding tiled dimensions into untiled dimensions is not supp

Error message

Folding tiled dimensions into untiled dimensions is not supported

What it means

During commutation, a group of source dimensions that must collapse into one reshaped dimension includes more tiled (untiled) dimensions than remain in the tiling stack — i.e. the reshape would fold tiled dimensions into an untiled dimension, which the algebra cannot represent.

Source

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

        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(
            "Folding tiled dimensions into untiled dimensions is not supported"
        )
      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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restructure the reshape so only untiled dims fold together, or only tiled dims with unit tiling
2. Apply the reshape to the untiled (materialized) array instead of the transformed ref
3. Reorder transforms so the untile happens before the reshape

Example fix

// before
flat = tiled_ref.reshape((16,))

// after
flat = jnp.asarray(tiled_ref).reshape((16,))

try: commute/reshape\nexcept NotImplementedError: materialize then reshape

When it happens

Trigger: A fold-only reshape that merges dimensions where tiling dims are exhausted, e.g. tiling (2,) on shape (4, 4) reshaped to (16,) with more dims to fold than tiled dims available.

Common situations: Flattening multi-dimensional tiled blocks in pallas kernels (reshape(-1)) where the minor dimension is tiled but the major ones are not.

Related errors

• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unfolding dimensions is not supported when commuting an `Un

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