{"record":{"id":"9f74312bde1472d6","repo":"jax-ml/jax","slug":"reshape-shape-shape-is-not-divisible-by-swizzle","errorCode":null,"errorMessage":"Reshape shape {shape} is not divisible by swizzle elements {self.swizzle_elems(aval.dtype)}","messagePattern":"Reshape shape (.+?) is not divisible by swizzle elements (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1229,"sourceCode":"  def swizzle_elems(self, dtype: jax.typing.DTypeLike | ir.Type) -> int:\n    if not isinstance(dtype, ir.Type):\n      dtype = mgpu_utils.dtype_to_ir_type(dtype)\n    return (self.swizzle * 8) // mgpu.bitwidth(dtype)\n\n  def commute_transpose(\n      self, _: jax_core.AbstractValue, transpose: state_types.TransposeTransform\n  ) -> tuple[state_types.TransposeTransform, UnswizzleRef]:\n    perm = transpose.permutation\n    if perm[-1] != len(perm) - 1:\n      raise ValueError(\"Can't transpose the swizzled dimension.\")\n    return transpose, self\n\n  def commute_reshape(\n      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform\n  ) -> tuple[state_types.ReshapeTransform, UnswizzleRef]:\n    shape = aval.shape\n    if shape[-1] != self.swizzle_elems(aval.dtype):\n      raise ValueError(\n          f\"Reshape shape {shape} is not divisible by swizzle elements\"\n          f\" {self.swizzle_elems(aval.dtype)}\"\n      )\n    return transform, self\n\n  def commute_ndindexer(\n      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer\n  ) -> tuple[indexing.NDIndexer, UnswizzleRef]:\n    if not hasattr(aval, \"dtype\"):\n      raise ValueError(\n          f\"Cannot commute unswizzle and indexer with {aval}, which does not\"\n          \" have a dtype\"\n      )\n    dtype = aval.dtype\n    swizzle_elems = self.swizzle_elems(dtype)\n    idxs = indexer.indices\n    if not idxs:\n      return indexer, self","sourceCodeStart":1211,"sourceCodeEnd":1247,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1211-L1247","documentation":"UnswizzleRef.commute_reshape requires that after a reshape, the minormost dimension of the new shape equals exactly the swizzle element count for the dtype (shape[-1] == swizzle_elems). A reshape to any other trailing extent would break the swizzled 128-bit vector grouping, so it is rejected.","triggerScenarios":"Calling reshape on a ref with a SwizzleTransform where the resulting last dimension is not swizzle_elems(dtype) (e.g. reshaping to trailing dim 64 or 256 when swizzle_elems is 128 for the dtype).","commonSituations":"Reshaping WGMMA operand/accumulator blocks in Mosaic kernels to feed another matmul; dtype changes that alter swizzle_elems (e.g. switching f32 to bf16) without adjusting the reshape target.","solutions":["Compute swizzle_elems for your dtype (swizzle_bytes*8 // dtype bit-width) and reshape so the last dim equals it","Unswizzle the ref first, reshape freely, then re-swizzle","Keep the trailing dimension fixed and fold size changes into leading dims instead"],"exampleFix":"// before\nref2 = ref.reshape((m, n))  # n != swizzle_elems -> ValueError\n// after\nse = swizzle_elems(dtype)\nref2 = ref.reshape((m * n // se, se))","handlingStrategy":"validation","validationCode":"se = swizzle_elems(dtype)\nassert new_shape[-1] == se, f'last dim must be {se}, got {new_shape[-1]}'","typeGuard":null,"tryCatchPattern":"try:\n    ref.reshape(new_shape)\nexcept ValueError:\n    unswizzle(ref).reshape(new_shape)","preventionTips":["Compute swizzle_elems from the dtype before choosing reshape targets","Fold size changes into leading dims, keep trailing dim = swizzle_elems"],"tags":["jax","pallas","reshape","swizzle","mosaic-gpu"],"backgroundTag":"layout-incompatible-reshape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}