{"record":{"id":"14e5aca59a8d06c1","repo":"jax-ml/jax","slug":"every-block-dimension-must-be-either-a-multiple-or","errorCode":null,"errorMessage":"Every block dimension must be either a multiple or factor of input. Got block {block_shape} for input {aval_in.shape}","messagePattern":"Every block dimension must be either a multiple or factor of input\\. Got block (.+?) for input (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":2143,"sourceCode":"    reps: tuple[int, ...],\n):\n  del reps\n  block_shape = block_transform.block_shape\n  aval_in = ctx.avals_in[0]\n  assert isinstance(aval_in, core.ShapedArray)\n  assert len(block_shape) == len(aval_in.shape)\n  if not all(isinstance(dim, (int, pallas_core.Squeezed))\n             for dim in block_shape):\n    raise NotImplementedError(\n        'tile with non-int block dimensions not supported yet'\n    )\n\n  if not all(\n      (pallas_core.get_block_size(block_dim) % in_dim == 0) or\n      (in_dim % pallas_core.get_block_size(block_dim) == 0)\n      for block_dim, in_dim in zip(block_shape, aval_in.shape)\n  ):\n    raise NotImplementedError(\n        'Every block dimension must be either a multiple or factor of input. '\n        f'Got block {block_shape} for input {aval_in.shape}'\n    )\n\n  new_shape = tuple(\n      block_dim if isinstance(block_dim, pallas_core.Squeezed)\n      else min(block_dim, in_dim)\n      for block_dim, in_dim in zip(block_shape, aval_in.shape)\n  )\n\n  def new_block_index_transform(*idxs):\n    original_idxs = block_transform.block_index_transform(*idxs)\n    return tuple(\n        0 if pallas_core.get_block_size(block_dim) >= in_dim\n        else orig_idx % (in_dim // pallas_core.get_block_size(block_dim))\n        for orig_idx, block_dim, in_dim in zip(\n            original_idxs, block_shape, aval_in.shape\n        )","sourceCodeStart":2125,"sourceCodeEnd":2161,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L2125-L2161","documentation":"Companion check to 1886: for tiling, every block dimension's effective size (via pallas_core.get_block_size) must be either an exact multiple or an exact factor of the corresponding input dimension. Anything else (e.g. block 6 vs input 4) can't be tiled whole-number-of-times, so the rule fails.","triggerScenarios":"Tile/broadcast usage inside a fused Pallas kernel where for some axis neither block_size % in_dim == 0 nor in_dim % block_size == 0 holds (Squeezed dims contribute size 1, which always divides).","commonSituations":"Block sizes chosen for hardware tiling (e.g. 128-wide tiles) applied to operands whose extents share no divisor relationship; partial reshapes leaving odd extents; changing block sizes without re-checking operand shapes.","solutions":["Choose block sizes that are exact multiples of the input dims on tiled axes (or exact factors)","Reshape/pad the input so each axis divides or is divided by the block size","Skip tiling on that axis by using a block dim equal to the input dim"],"exampleFix":"// before\nblock_shape = (6, 128)  # input shape (4, 128): 6 % 4 != 0 and 4 % 6 != 0\n// after\nblock_shape = (8, 128)  # 8 % 4 == 0\n","handlingStrategy":"validation","validationCode":"from jax._src.pallas import core as pc\nassert all((pc.get_block_size(b) % i == 0) or (i % pc.get_block_size(b) == 0)\n           for b, i in zip(block_shape, x.shape)), 'block dims must be multiple/factor of input'","typeGuard":"def block_divides_or_divisible(block_shape, in_shape) -> bool:\n    return all((b % i == 0) or (i % b == 0) for b, i in zip(block_shape, in_shape))","tryCatchPattern":"try:\n    out = fused_tile_fn(x)\nexcept NotImplementedError as e:\n    if 'multiple or factor' in str(e):\n        out = unfused_tile(x)\n    else:\n        raise","preventionTips":["Pick hardware tile sizes that are powers of 2 dividing/ multiples of operand extents","Reshape operands so extents divide block sizes","Validate block/operand divisibility in a pre-launch assert"],"tags":["jax","pallas","tile","divisibility","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}