{"record":{"id":"007906e38e4b91cd","repo":"jax-ml/jax","slug":"expected-tiling-to-be-at-most-rank-of-shape-got-t","errorCode":null,"errorMessage":"Expected tiling to be at most rank of shape. Got tiling: {tiling} (rank: {len(tiling)}) and shape {shape} (rank: {len(shape)}).","messagePattern":"Expected tiling to be at most rank of shape\\. Got tiling: (.+?) \\(rank: (.+?)\\) and shape (.+?) \\(rank: (.+?)\\)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":1784,"sourceCode":"\n  def refine(\n      self,\n      *,\n      chunk: ir.Value | None = None,\n      num_chunks: int | None = None,\n      chunk_size: int | None = None,\n  ):\n    return Partition1D(\n        self.partition.target_block_shape[0],\n        num_chunks=num_chunks,\n        chunk_size=chunk_size,\n        base_offset=self.get_base(chunk) if chunk is not None else None,\n    )\n\n\ndef tile_shape(shape, tiling):\n  if len(tiling) > len(shape):\n    raise ValueError(\n        \"Expected tiling to be at most rank of shape. Got tiling:\"\n        f\" {tiling} (rank: {len(tiling)}) and shape {shape} (rank:\"\n        f\" {len(shape)}).\"\n    )\n  if not tiling:\n    return shape\n  tiling_rank = len(tiling)\n  for s, t in zip(shape[-tiling_rank:], tiling):\n    if s % t:\n      raise ValueError(\"Non-divisible tiling:\", shape, tiling)\n  return (\n      *shape[:-tiling_rank],\n      *(s // t for s, t in zip(shape[-tiling_rank:], tiling)),\n      *tiling,\n  )\n\n\ndef warp_tree_reduce(value, op, group_size):","sourceCodeStart":1766,"sourceCodeEnd":1802,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L1766-L1802","documentation":"tile_shape tiles the trailing dimensions of a shape with a given tiling; the tiling's rank (its length) must not exceed the shape's rank. If tiling has more dimensions than shape, the per-dimension zip would silently mis-align, so the helper raises immediately.","triggerScenarios":"Calling tile_shape(shape=(256,), tiling=(8, 8)) or any call site where the tiling tuple/list is longer than the shape tuple. Common in kernels configured with a 2D warp tile applied to 1D buffers.","commonSituations":"Copy-pasting a 2D matmul tiling config into a 1D reduction/normalization kernel; changing tensor rank without updating the tiling constants; mismatch between layout rank and block tiling in a Mosaic kernel.","solutions":["Reduce the tiling to at most the rank of shape (drop leading tiling dims or align them to trailing dims)","Increase the shape rank if the tensor was meant to be multi-dimensional","Check where the tiling was constructed (e.g., a block spec) and fix the rank mismatch at the source"],"exampleFix":"# before\ntile_shape(shape=(4096,), tiling=(8, 128))\n# after\ntile_shape(shape=(4096,), tiling=(128,))","handlingStrategy":"validation","validationCode":"assert len(tiling) <= len(shape), f'tiling rank {len(tiling)} > shape rank {len(shape)}'\nresult = tile_shape(shape, tiling)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive tiling length from the shape's rank rather than hardcoding","Assert rank consistency in kernel setup tests"],"tags":["mosaic-gpu","tiling","shape-validation"],"backgroundTag":"rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}