{"record":{"id":"f40a0aa6fdc1ff9a","repo":"jax-ml/jax","slug":"the-pallas-tpu-lowering-currently-requires-that-ra","errorCode":null,"errorMessage":"The Pallas TPU lowering currently requires that rank 1 block shapes, either 1) the first (and only) dimension of the block shape is equal to the first (and only) dimension of the array shape, or 2) the first (and only) dimension of the block shape is a multiple of {chunk_size}, or 3) the first (and only) dimension of the block shape is a power of 2 and at least the tiling size ({min_tiling} = 128 * (32 // {dtypes.itemsize_bits(physical_dtype)})) of the array shape. {err_details}","messagePattern":"The Pallas TPU lowering currently requires that rank 1 block shapes, either 1\\) the first \\(and only\\) dimension of the block shape is equal to the first \\(and only\\) dimension of the array shape, or 2\\) the first \\(and only\\) dimension of the block shape is a multiple of (.+?), or 3\\) the first \\(and only\\) dimension of the block shape is a power of 2 and at least the tiling size \\((.+?) = 128 \\* \\(32 // (.+?)\\)\\) of the array shape\\. (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":1037,"sourceCode":"        )\n    else:\n      assert rank == 1\n      if bm.array_aval.dtype == jnp.bool_:\n        bitwidth = dtypes.itemsize_bits(BOOL_MEMREF_TYPE)\n      else:\n        bitwidth = dtypes.itemsize_bits(physical_dtype)\n      packing = 32 // bitwidth\n      sublane_count = tpu_info.get_tpu_info().num_sublanes\n      lane_count = tpu_info.get_tpu_info().num_lanes\n      min_tiling = lane_count * packing\n      chunk_size = sublane_count * lane_count\n      feasible_block_size = (\n          bs0 == as0\n          or bs0 % chunk_size == 0\n          or (bs0 >= min_tiling and (bs0 & (bs0 - 1)) == 0)  # power of 2\n      )\n      if not feasible_block_size:\n        raise ValueError(\n            \"The Pallas TPU lowering currently requires that rank 1 block\"\n            \" shapes, either 1) the first (and only) dimension of the block\"\n            \" shape is equal to the first (and only) dimension of the array\"\n            \" shape, or 2) the first (and only) dimension of the block shape\"\n            f\" is a multiple of {chunk_size}, or 3) the first (and only)\"\n            \" dimension of the block shape is a power of 2 and at least the\"\n            f\" tiling size ({min_tiling} = 128 * (32 //\"\n            f\" {dtypes.itemsize_bits(physical_dtype)})) of the array shape. \"\n            + err_details()\n        )\n\n\ndef lower_jaxpr_to_pipelined_module(\n    lowering_context: mlir.LoweringRuleContext,\n    grid_mapping: pallas_core.GridMapping,\n    jaxpr: jax_core.Jaxpr,\n    *,\n    dimension_semantics: Sequence[tpu_core.DimensionSemantics] | None,","sourceCodeStart":1019,"sourceCodeEnd":1055,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L1019-L1055","documentation":"For rank-1 arrays on TPU, the Pallas lowering can only pipeline certain block sizes: the block size must equal the array length, be a multiple of chunk_size, or be a power of 2 that is at least min_tiling = 128 * (32 // itemsize_bits(dtype)). This is because 1-D access must map efficiently onto the TPU's rectangular VLM layouts.","triggerScenarios":"A rank-1 BlockSpec on TPU with a block size that is not equal to the array size, not a multiple of chunk_size, and not a power-of-2 >= min_tiling (e.g. a 100-element block of a 1000-element float32 array).","commonSituations":"1-D kernels (e.g. elementwise maps over vectors, sorts, scans) tiled with odd block sizes; low-precision dtypes (f8/bf16) raising min_tiling to 512; block sizes chosen to match a dataset size rather than hardware constraints.","solutions":["Use a power-of-2 block size at least the min_tiling for your dtype (e.g. >=128 for 32-bit types, >=512 for 8-bit types)","Use a block size that is a multiple of chunk_size (chunk_size = 128 * 128 // itemsize_bits)","Set the block size equal to the full array length (with a grid of 1) if it fits in memory","Switch the operand to 2-D by adding a trailing dimension of 1 if the algorithm permits"],"exampleFix":"# before\nBlockSpec(block_shape=(100,), index_map=lambda i: i)  # 100 not feasible\n\n# after\nBlockSpec(block_shape=(128,), index_map=lambda i: i)  # power of 2 >= 128","handlingStrategy":"validation","validationCode":"import numpy as np\ndef feasible_rank1(bs0, as0, dtype):\n    import jax.numpy as jnp\n    itemsize_bits = jnp.dtype(dtype).itemsize * 8\n    chunk = 128 * 128 // itemsize_bits\n    min_tiling = 128 * (32 // itemsize_bits)\n    return bs0 == as0 or bs0 % chunk == 0 or (bs0 >= min_tiling and (bs0 & (bs0-1)) == 0)","typeGuard":"def is_valid_1d_block(bs0, as0, chunk_size, min_tiling) -> bool:\n    return bs0 == as0 or bs0 % chunk_size == 0 or \\\n           (bs0 >= min_tiling and bs0 & (bs0 - 1) == 0)","tryCatchPattern":"try:\n    pallas_call(kernel, out, grid=grid)\nexcept ValueError as e:\n    if 'rank 1 block' in str(e):\n        out = pallas_call(kernel, reshape_2d_spec_out, grid=grid)  # reshape to 2D and retry","preventionTips":["For rank-1 kernels use power-of-2 blocks >= 128 (>= 512 for 8-bit dtypes)","Remember min_tiling doubles for 16-bit and quadruples for 8-bit dtypes","Prefer 2-D shapes with a trailing dim of 1 for 1-D algorithms"],"tags":["jax","pallas","tpu","block-shape","rank-1"],"backgroundTag":"pallas-tpu-block-alignment-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}