{"record":{"id":"f2bf6bdfa20b9d0f","repo":"jax-ml/jax","slug":"the-pallas-tpu-lowering-currently-requires-that-th","errorCode":null,"errorMessage":"The Pallas TPU lowering currently requires that the last two dimensions of your block shape are divisible by 8 and 128 respectively, or be equal to the respective dimensions of the overall array. {extra_msg}{err_details}","messagePattern":"The Pallas TPU lowering currently requires that the last two dimensions of your block shape are divisible by 8 and 128 respectively, or be equal to the respective dimensions of the overall array\\. (.+?)(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":1012,"sourceCode":"      bs1, as1 = unmapped_bs[-2], physical_array_shape[-2]\n    else:\n      bs1, as1 = 1, 1\n\n    if rank >= 2:\n      evenly_divisible = (\n          (bs0 == as0 or bs0 % 128 == 0) and\n          (bs1 == as1 or bs1 % 8 == 0)\n      )\n      if not evenly_divisible:\n        extra_msg = \"\"\n        if pallas_core.dynamic_shapes_export_enabled():\n          extra_msg = (\n              \" In dynamic shape export - your kernel symbolic args must be\"\n              \" annotated with constraints where the computation *after*\"\n              \" applying any grid mapping is divisible by 8 and 128\"\n              \" respectively. Ex: (mod(floordiv(m_dim, grid_size), 8) == 0))\"\n          )\n        raise ValueError(\n            \"The Pallas TPU lowering currently requires that the last two \"\n            \"dimensions of your block shape are divisible by 8 and 128 \"\n            \"respectively, or be equal to the respective dimensions of the \"\n            \"overall array. \"\n            + extra_msg\n            + err_details()\n        )\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","sourceCodeStart":994,"sourceCodeEnd":1030,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L994-L1030","documentation":"On TPU, the vector/VLM units require the last two dimensions of each block shape to be divisible by 8 and 128 respectively, unless they equal the corresponding full-array dimensions. The Pallas TPU lowering enforces this hardware alignment constraint before emitting the pipelined Mosaic module, and raises this ValueError (with extra guidance for dynamic shape export) when violated.","triggerScenarios":"A pallas_call on TPU whose BlockSpec block_shape has a last dim not divisible by 128 (and not equal to the array's last dim), or second-to-last dim not divisible by 8; with dynamic shape export, symbolic dims lacking constraints like (mod(floordiv(m_dim, grid_size), 8) == 0).","commonSituations":"Porting GPU Pallas/triton kernels to TPU with arbitrary tile sizes (e.g. 100x100 blocks); using non-multiple-of-128 tile sizes for the fastest dimension; dynamic shape export where symbolic dimensions aren't annotated with divisibility constraints.","solutions":["Change the last dimension of your block shape to a multiple of 128 (e.g. 128, 256) or to the full array dimension","Make the second-to-last block dimension a multiple of 8 or equal to the array's dimension","For dynamic shape export, add divisibility constraints to symbolic args, e.g. (mod(floordiv(m_dim, grid_size), 8) == 0) and ... 128 == 0","Pad the input arrays so the block constraints can be satisfied and slice the result afterwards"],"exampleFix":"# before\nBlockSpec(block_shape=(1024, 100), index_map=...)\n\n# after\nBlockSpec(block_shape=(1024, 128), index_map=...)  # last dim multiple of 128","handlingStrategy":"validation","validationCode":"def check_tpu_block_alignment(block_shape, array_shape):\n    ok_last = block_shape[-1] % 128 == 0 or block_shape[-1] == array_shape[-1]\n    ok_prev = block_shape[-2] % 8 == 0 or block_shape[-2] == array_shape[-2]\n    if not (ok_last and ok_prev):\n        raise ValueError('last two block dims must be div by 128/8 or equal array dims')","typeGuard":"def is_tpu_aligned(block, arr) -> bool:\n    return (block[-1] % 128 == 0 or block[-1] == arr[-1]) and \\\n           (block[-2] % 8 == 0 or block[-2] == arr[-2])","tryCatchPattern":"try:\n    pallas_call(kernel, out_specs, in_specs, grid=grid)\nexcept ValueError as e:\n    if 'divisible by 8 and 128' in str(e):\n        block = next_power_of_2_or_128(block)  # adjust tile and retry","preventionTips":["Default to tile sizes whose last dim is a multiple of 128 and second-to-last a multiple of 8","For dynamic shape export, add divisibility constraints on symbolic dims","Pad inputs to friendly sizes and slice outputs"],"tags":["jax","pallas","tpu","alignment","dynamic-shapes"],"backgroundTag":"pallas-tpu-block-alignment-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}