jax-ml/jax · error · NotImplementedError

Commuting a `UntilingTransform` with a `ReshapeTransform`…

Error message

Commuting a `UntilingTransform` with a `ReshapeTransform` is not supported when the target shape has 0 dimensions

What it means

Raised when trying to commute (reorder) an UntilingTransform past a ReshapeTransform whose target shape is zero-dimensional (e.g. reshaping to ()). The Mosaic GPU transform algebra cannot express untiled semantics for an empty shape, so the pass fails loudly rather than producing wrong indexing.

Solutions

Avoid reshaping the block to a 0-d shape; keep at least one dimension of size 1 (reshape to (1,) instead of ())
Apply the untile/reshape in a different order or materialize the value into a normal array before reshaping
Restructure the kernel so the tiled reference is fully consumed before any scalar reshape

Example fix

// before
block = block.reshape(())

// after
block = block.reshape((1,))

Defensive patterns

Strategy: validation

Validate before calling

if not isinstance(shape, tuple) or len(shape) == 0: raise ValueError('reshape target must have >=1 dim for tiled refs')

Prevention

Never reshape pallas blocks to 0-d; use (1,)
Keep at least one dimension on reduced blocks in kernels

When it happens

Trigger: Calling commute_reshape on an UntilingTransform where transform.shape is empty — i.e. a pallas kernel path where a reference with tiling transforms is reshaped to a scalar shape ().

Common situations: Pipelines that reduce a block to a scalar (e.g. after a full reduction) and then use the result through a transform stack; unexpected empty shapes from degenerate grid dims or size-1 arrays reshaped to ().

Related errors

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

Appendix: source

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

          idxs_after_tiling.append(indexing.Slice(new_start, new_size))
        case _:
          raise TypeError(f"Unsupported index type: {type(idx)}")
    assert all(a % b == 0 for a, b in zip(untiled_shape, self.tiling))
    tiled_shape = [
        *(a // b for a, b in zip(untiled_shape, self.tiling)),
        *self.tiling,
    ]
    new_indexer = indexing.NDIndexer.from_indices_shape(
        indices=(*untiled_idxs, *idxs_after_tiling),
        shape=(*leading_shape, *tiled_shape)
    )
    return new_indexer, self

  def commute_reshape(
      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform
  ) -> tuple[state_types.ReshapeTransform, UntilingTransform]:
    if not transform.shape:
      raise NotImplementedError(
          "Commuting a `UntilingTransform` with a `ReshapeTransform` is not "
          "supported when the target shape has 0 dimensions"
      )
    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])

View on GitHub (pinned to 1e1c6a8fc0)