{"record":{"id":"33fd0439348905af","repo":"jax-ml/jax","slug":"tile-with-non-int-block-dimensions-not-supported-y","errorCode":null,"errorMessage":"tile with non-int block dimensions not supported yet","messagePattern":"tile with non-int block dimensions not supported yet","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":2103,"sourceCode":"\n\n@register_eval_rule(lax.split_p)\ndef _split_eval_rule(\n    eval_ctx: KernelEvalContext, x, sizes: Sequence[int], axis: int\n):\n  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(","sourceCodeStart":2085,"sourceCodeEnd":2121,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L2085-L2121","documentation":"Raised by the fuser's tile operation evaluation rule: tiling (repeating an input to fill a larger output block) requires every non-squeezed block dimension to be a concrete int. A block dim of None (unbounded/undefined at trace time) makes the repeat count uncomputable, so it's rejected.","triggerScenarios":"Using jnp.tile / lax.tile inside a fused Pallas region where the output BlockSpec has a None (unbounded) block dimension along the tiled axis (after filtering out Squeezed dims).","commonSituations":"Hand-written BlockSpecs with None dims for flexible sequence lengths combined with broadcasting/tiling; kernels designed for ragged or dynamic shapes that then apply tile; version changes where tile fusion was added with this restriction.","solutions":["Set explicit integer block sizes on the output BlockSpec for the tiled axes","Replace tile with broadcast_to + reshape when the target shape is static","Move the tile outside the kernel and pass the pre-tiled array as input"],"exampleFix":"// before\nspec = BlockSpec((None, 64), ...)  # unbounded dim\nout = fuse(jnp.tile, ...)(x, (3, 1))\n// after\nspec = BlockSpec((x.shape[0] * 3, 64), ...)  # concrete block dim\n","handlingStrategy":"validation","validationCode":"eff = [d for d in spec.block_shape if not isinstance(d, pallas_core.Squeezed)]\nassert all(isinstance(d, int) for d in eff), 'tile needs concrete int block dims'","typeGuard":"def tile_dims_concrete(block_shape) -> bool:\n    return all(isinstance(d, int) for d in block_shape if not isinstance(d, pallas_core.Squeezed))","tryCatchPattern":"try:\n    y = fused_tile(x, reps)\nexcept NotImplementedError as e:\n    if 'non-int block dimensions' in str(e):\n        y = jnp.tile(x, reps)\n    else:\n        raise","preventionTips":["Avoid None block dims on axes you tile","Prefer broadcast_to+reshape for static shapes","Keep tiled-axis block sizes explicit in BlockSpec"],"tags":["jax","pallas","tile","block-spec","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}