jax-ml/jax · error · ValueError

Shape along concat dimension {dimension} must be divisible b

Error message

Shape along concat dimension {dimension} must be divisible by the block shape {block_shape[dimension]} for all children. Got shape {aval.shape}.

What it means

Raised by _concatenate_eval_rule when a child array's extent along the concatenation dimension is not a multiple of the block size on that dimension. The fuser decomposes concat into per-child block ranges, which requires each child to occupy an integer number of blocks.

Source

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

  is_element_block = [isinstance(bd, pallas_core.Element) for bd in block_shape]
  if any(is_element_block):
    raise NotImplementedError(
        'Concatenation with Element indexing is not yet supported.'
    )
  block_dim = block_shape[dimension]
  if block_dim is None:
    block_dim = 1

  if block_dim == sum(aval.shape[dimension] for aval in ctx.avals_in):
    # Handle special case if the block contains all of the concatenated
    # array.
    return jax.lax.concatenate(args, dimension=dimension)

  num_blocks = []
  for aval in ctx.avals_in:
    assert isinstance(aval, core.ShapedArray)
    if aval.shape[dimension] % block_dim != 0:
      raise ValueError(
          f'Shape along concat dimension {dimension} must be divisible by the'
          f' block shape {block_shape[dimension]} for all children. Got shape'
          f' {aval.shape}.'
      )
    num_blocks.append(aval.shape[dimension] // block_dim)
  ends = np.cumsum(num_blocks).astype(np.int32)
  starts = np.concatenate(([0], ends[:-1])).astype(np.int32)

  block_indices = ctx.get_out_block_indices()[0]
  block_idx = block_indices[dimension]
  valid_index = 0
  for i in range(len(ctx.avals_in)):
    start, end = starts[i], ends[i]
    is_valid = (start <= block_idx) & (block_idx < end)
    valid_index = jax.lax.select(is_valid, i, valid_index)
  out_dtype = args[0].dtype
  args = [a.astype(jnp.float32) if a.dtype == jnp.bfloat16 else a for a in args]
  valid_block = jax.lax.select_n(valid_index, *args)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Choose a block size along the concat dimension that divides every child's extent (e.g. gcd of the child sizes)
Pad children to multiples of the block size before concatenation (and trim afterwards if needed)
If the block covers the entire concatenated extent, ensure the special whole-block path applies (block_dim == total extent)
Restructure to avoid in-kernel concat of ragged children

Example fix

# before
# children of length 10 and 6, block size 4 -> 10 % 4 != 0
out[...] = jnp.concatenate([x, y], axis=0)

# after
# block size 2 divides both 10 and 6
BlockSpec(block_shape=(2, ...), ...)
out[...] = jnp.concatenate([x, y], axis=0)

Defensive patterns

Strategy: validation

Validate before calling

import math
block_dim = block_shape[dimension]
assert all(a.shape[dimension] % block_dim == 0 for a in arrays), \
    'each child extent must be divisible by the concat-axis block size'

Try / catch

try:
    kernel = pallas_call(fn, out_spec=BlockSpec(block_shape=(bs, ...)))(
        jnp.concatenate, ...)  #示意
except ValueError as e:
    if 'divisible by the block shape' in str(e):
        bs = math.gcd(*(a.shape[dimension] for a in arrays))
        # retry with gcd block size
    else:
        raise

Prevention

Pick block size = gcd of child extents along the concat axis
Pad children to block multiples before concat
Add a divisibility assert before launching the kernel

When it happens

Trigger: jnp.concatenate where some input's shape[dimension] % block_shape[dimension] != 0, e.g. concatenating arrays of length 10 with block size 4 along that axis in a pallas kernel.

Common situations: Concatenating arrays with remainder sizes after choosing block sizes that only divide the total; changing block sizes during tuning; concatenating differently-sized buffers.

Related errors

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