{"record":{"id":"3d82e2b0948a9d58","repo":"jax-ml/jax","slug":"unsupported-shape-x-shape","errorCode":null,"errorMessage":"Unsupported shape: {x.shape}","messagePattern":"Unsupported shape: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1316,"sourceCode":"\n@tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True)\nclass ExpandLeadingBatchDimensionsTransform(state_types.Transform):\n  \"\"\"The inverse of CollapseLeadingBatchDimensionsTransform.\n\n  Specifically, it maps `(m, math.prod(batch_shape) * n)` to `(*batch_shape, m,\n  n)`.\n  \"\"\"\n\n  batch_shape: tuple[int, ...] = jax.tree.static()\n\n  def transform_type(\n      self, x: jax_core.AbstractValue\n  ) -> state_types.AbstractRef:\n    match x:\n      case jax_core.ShapedArray():\n        if x.ndim != 2:\n          raise ValueError(f\"Unsupported shape: {x.shape}\")\n        batch_size = math.prod(self.batch_shape)\n        if x.shape[1] % batch_size != 0:\n          raise ValueError(\n              f\"Second dimension {x.shape[1]} must be divisible by batch_size\"\n              f\" {batch_size}\"\n          )\n        transformed_shape = self.batch_shape + (\n            x.shape[0],\n            x.shape[1] // batch_size,\n        )\n        return x.update(shape=transformed_shape)\n      case state_types.AbstractRef():\n        return x.update(inner_aval=self.transform_type(x.inner_aval))\n      case _:\n        raise TypeError(f\"Unsupported type: {x}\")\n\n  def commute_ndindexer(\n      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer","sourceCodeStart":1298,"sourceCodeEnd":1334,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1298-L1334","documentation":"The inverse transform (undoing batch-dim expansion) requires the physical aval to be exactly 2-D, since the layout maps a (batch..., m, n) logical shape to a flat (m, batch*n) physical shape. ndim != 2 raises ValueError('Unsupported shape').","triggerScenarios":"Applying the batch-expansion undo/transform to a ShapedArray with ndim other than 2 — e.g. trying to interpret a 3-D or 1-D physical buffer as an expanded-batch layout.","commonSituations":"Mismatch between the shape of a physical buffer and the declared batch_shape in a Mosaic kernel; hand-constructed avals passed to get_ref_aval/to_block_mapping.","solutions":["Make the physical aval 2-D (m, batch*n) before the transform","Verify batch_shape passed to the transform matches how the buffer was allocated","Reconstruct the buffer with the correct 2-D physical layout"],"exampleFix":"// before\nphys = ShapedArray((b, m, n), dt)  # ndim 3 -> ValueError\n// after\nphys = ShapedArray((m, b * n), dt)","handlingStrategy":"validation","validationCode":"assert x.ndim == 2, f'physical aval must be 2-D, got shape {x.shape}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Allocate physical buffers as (m, batch*n) 2-D arrays","Keep declared batch_shape consistent with the buffer layout"],"tags":["jax","pallas","shape-validation","ndim"],"backgroundTag":"invalid-shape-for-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}