jax-ml/jax · error · NotImplementedError

tile with non-int block dimensions not supported yet

Error message

tile with non-int block dimensions not supported yet

What it means

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.

Source

Thrown at jax/_src/pallas/fuser/block_spec.py:2103

@register_eval_rule(lax.split_p)
def _split_eval_rule(
    eval_ctx: KernelEvalContext, x, sizes: Sequence[int], axis: int
):
  del eval_ctx
  return lax.split(x, sizes=sizes, axis=axis)


@register_eval_rule(lax.tile_p)
def _tile_eval_rule(
    eval_ctx: KernelEvalContext, x, reps: tuple[int, ...]
):
  block_spec = eval_ctx.out_block_specs[0]
  block_shape = tuple(d for d in block_spec.block_shape
                      if not isinstance(d, pallas_core.Squeezed))
  if not all(isinstance(dim, int) for dim in block_shape):
    raise NotImplementedError(
        'tile with non-int block dimensions not supported yet'
    )
  if not all(
      out_dim % in_dim == 0 for out_dim, in_dim in zip(block_shape, x.shape)
  ):
    raise NotImplementedError(
        'Block size must be a multiple of the input size. '
        f'Got block {block_shape=} but input {x.shape}.'
    )
  reps_in_block = [
      out_dim // in_dim if out_dim >= in_dim else 1
      for out_dim, in_dim in zip(block_shape, x.shape)
  ]
  return lax.tile(x, reps_in_block)


@register_pull_block_spec_rule(lax.tile_p)
def _tile_pull_rule(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set explicit integer block sizes on the output BlockSpec for the tiled axes
2. Replace tile with broadcast_to + reshape when the target shape is static
3. Move the tile outside the kernel and pass the pre-tiled array as input

Example fix

// before
spec = BlockSpec((None, 64), ...)  # unbounded dim
out = fuse(jnp.tile, ...)(x, (3, 1))
// after
spec = BlockSpec((x.shape[0] * 3, 64), ...)  # concrete block dim

eff = [d for d in spec.block_shape if not isinstance(d, pallas_core.Squeezed)]
assert all(isinstance(d, int) for d in eff), 'tile needs concrete int block dims'

def tile_dims_concrete(block_shape) -> bool:
    return all(isinstance(d, int) for d in block_shape if not isinstance(d, pallas_core.Squeezed))

try:
    y = fused_tile(x, reps)
except NotImplementedError as e:
    if 'non-int block dimensions' in str(e):
        y = jnp.tile(x, reps)
    else:
        raise

When it happens

Trigger: 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).

Common situations: 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.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton
• Unsupported block dimension type: {type(dim)}. Allowed types
• Unsupported block shape type: {type(dim)}
• Array type must have a `shape` attribute, but got {type(arra

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/33fd0439348905af. Report an issue: GitHub.