{"record":{"id":"0a30c8b6bf8f45a4","repo":"jax-ml/jax","slug":"block-size-must-be-a-multiple-of-the-input-size-g","errorCode":null,"errorMessage":"Block size must be a multiple of the input size. Got block {block_shape=} but input {x.shape}.","messagePattern":"Block size must be a multiple of the input size\\. Got block (.+?) but input (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":2109,"sourceCode":"  del eval_ctx\n  return lax.split(x, sizes=sizes, axis=axis)\n\n\n@register_eval_rule(lax.tile_p)\ndef _tile_eval_rule(\n    eval_ctx: KernelEvalContext, x, reps: tuple[int, ...]\n):\n  block_spec = eval_ctx.out_block_specs[0]\n  block_shape = tuple(d for d in block_spec.block_shape\n                      if not isinstance(d, pallas_core.Squeezed))\n  if not all(isinstance(dim, int) for dim in block_shape):\n    raise NotImplementedError(\n        'tile with non-int block dimensions not supported yet'\n    )\n  if not all(\n      out_dim % in_dim == 0 for out_dim, in_dim in zip(block_shape, x.shape)\n  ):\n    raise NotImplementedError(\n        'Block size must be a multiple of the input size. '\n        f'Got block {block_shape=} but input {x.shape}.'\n    )\n  reps_in_block = [\n      out_dim // in_dim if out_dim >= in_dim else 1\n      for out_dim, in_dim in zip(block_shape, x.shape)\n  ]\n  return lax.tile(x, reps_in_block)\n\n\n@register_pull_block_spec_rule(lax.tile_p)\ndef _tile_pull_rule(\n    ctx: PullRuleContext,\n    block_transform: BlockIndexTransform,\n    *,\n    reps: tuple[int, ...],\n):\n  del reps","sourceCodeStart":2091,"sourceCodeEnd":2127,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L2091-L2127","documentation":"The tile rule additionally requires each output block dimension to be an integer multiple of the corresponding input dimension. If out_dim % in_dim != 0 the repeat count is fractional, which cannot be expressed, so the fuser raises NotImplementedError.","triggerScenarios":"jnp.tile / lax.tile inside a fused Pallas region where the output block shape is not a whole multiple of the input shape along at least one axis, e.g. tiling a length-3 axis into a block of 8.","commonSituations":"Mismatched tiling factors between BlockSpec and the tile call; padding targets whose size isn't divisible by the input extent; combining squeeze/reshape with tile so shapes drift.","solutions":["Make the output block shape an exact multiple of the input shape on every axis (adjust block size or pad the input first)","Pad the input (e.g. to the next multiple) before tiling","Hoist the tile out of the fused region"],"exampleFix":"// before\ny = fuse(jnp.tile)(x, (8, 1))  # x.shape[0] == 3, 8 % 3 != 0\n// after\nx_padded = jnp.pad(x, ((0, 1), (0, 0)))  # now length 4\ny = fuse(jnp.tile)(x_padded, (8, 1))  # 8 % 4 == 0\n","handlingStrategy":"validation","validationCode":"assert all(b % i == 0 for b, i in zip(block_shape, x.shape) if i > 0), 'block must be multiple of input on every axis'","typeGuard":"def tile_multiple_ok(block_shape, in_shape) -> bool:\n    return all(b % i == 0 for b, i in zip(block_shape, in_shape))","tryCatchPattern":"try:\n    y = fused_tile(x, reps)\nexcept NotImplementedError as e:\n    if 'multiple of the input size' in str(e):\n        y = jnp.tile(x, reps)  # or pad input to a multiple first\n    else:\n        raise","preventionTips":["Pad inputs to whole multiples of the block before tiling","Compute reps from block/input ratios, never hardcode","Add shape assertions in tests"],"tags":["jax","pallas","tile","shape-mismatch","divisibility"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}