{"record":{"id":"36d2d4999b716ce8","repo":"jax-ml/jax","slug":"commuting-a-untilingtransform-with-a-reshapetra","errorCode":null,"errorMessage":"Commuting a `UntilingTransform` with a `ReshapeTransform` is not supported when the target shape has 0 dimensions","messagePattern":"Commuting a `UntilingTransform` with a `ReshapeTransform` is not supported when the target shape has 0 dimensions","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":856,"sourceCode":"          idxs_after_tiling.append(indexing.Slice(new_start, new_size))\n        case _:\n          raise TypeError(f\"Unsupported index type: {type(idx)}\")\n    assert all(a % b == 0 for a, b in zip(untiled_shape, self.tiling))\n    tiled_shape = [\n        *(a // b for a, b in zip(untiled_shape, self.tiling)),\n        *self.tiling,\n    ]\n    new_indexer = indexing.NDIndexer.from_indices_shape(\n        indices=(*untiled_idxs, *idxs_after_tiling),\n        shape=(*leading_shape, *tiled_shape)\n    )\n    return new_indexer, self\n\n  def commute_reshape(\n      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform\n  ) -> tuple[state_types.ReshapeTransform, UntilingTransform]:\n    if not transform.shape:\n      raise NotImplementedError(\n          \"Commuting a `UntilingTransform` with a `ReshapeTransform` is not \"\n          \"supported when the target shape has 0 dimensions\"\n      )\n    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])","sourceCodeStart":838,"sourceCodeEnd":874,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L838-L874","documentation":"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.","triggerScenarios":"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 ().","commonSituations":"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 ().","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"],"exampleFix":"// before\nblock = block.reshape(())\n\n// after\nblock = block.reshape((1,))","handlingStrategy":"validation","validationCode":"if not isinstance(shape, tuple) or len(shape) == 0: raise ValueError('reshape target must have >=1 dim for tiled refs')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never reshape pallas blocks to 0-d; use (1,)","Keep at least one dimension on reduced blocks in kernels"],"tags":["jax","pallas","mosaic-gpu","reshape","transforms"],"backgroundTag":"unsupported-transform-commutation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}