jax-ml/jax · error · NotImplementedError

Stacking only supported when the block size along the stack

Error message

Stacking only supported when the block size along the stack axis equals the number of inputs. Got block_dim={block_dim}, expected {n}.

What it means

Raised by JAX's Pallas kernel fuser when an operation attempts to stack N inputs along an axis, but the block shape's size along that stack axis does not equal the number of inputs. The fuser can only fuse a stacking pattern (e.g. lax.concatenate/stack lowered to a single block) when block_dim == len(ctx.avals_in). Any mismatch between the declared block shape and the operand count is rejected.

Source

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

def _stack_pull_rule(
    ctx: PullRuleContext,
    block_transform: BlockIndexTransform,
    *,
    axis: int,
):
  block_shape = block_transform.block_shape
  is_element_block = [isinstance(bd, pallas_core.Element) for bd in block_shape]
  if any(is_element_block):
    raise NotImplementedError(
        'Stack with Element indexing is not yet supported.'
    )
  block_dim = block_shape[axis]
  if block_dim is None or isinstance(block_dim, pallas_core.Squeezed):
    block_dim = 1

  n = len(ctx.avals_in)
  if block_dim != n:
    raise NotImplementedError(
        "Stacking only supported when the block size along the stack axis "
        f"equals the number of inputs. Got block_dim={block_dim}, expected {n}."
    )

  new_block_shape = list(block_transform.block_shape)
  new_block_shape.pop(axis)

  def make_block_transform(child_index: int):
    def new_block_index_transform(*idxs):
      idx = list(block_transform.block_index_transform(*idxs))
      idx.pop(axis)
      return tuple(idx)

    return block_transform.replace(
        block_shape=tuple(new_block_shape),
        block_index_transform=new_block_index_transform
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set block_shape[axis] to exactly the number of inputs being stacked (e.g. for jnp.stack([a, b], axis=0) use block size 2 on axis 0)
2. Split the stack/concatenate outside the fused kernel and pass inputs as separate block-mapped arguments
3. Check that the stack axis is not marked with a None (unbounded) or Squeezed block dim before fusion

Example fix

// before
spec = BlockSpec((None, 128, 128), index_map=lambda i, j, k: (i, j, k))
out = fuse(jnp.stack, ...)([a, b], axis=0)
// after
spec = BlockSpec((2, 128, 128), index_map=lambda i, j, k: (i, j, k))  # 2 == number of inputs

n = len(inputs)
assert spec.block_shape[axis] in (n,) or (spec.block_shape[axis] is None and n == 1), 'stack axis block size must equal number of inputs'

def stack_spec_ok(block_shape, axis, inputs) -> bool:
    bd = block_shape[axis]
    return bd is None and len(inputs) == 1 or bd == len(inputs)

try:
    out = fused_kernel(*inputs)
except NotImplementedError as e:
    if 'Stacking only supported' in str(e):
        out = jnp.stack(inputs, axis=axis)  # unfused fallback
    else:
        raise

When it happens

Trigger: Calling a Pallas kernel or jax fusion pass where multiple inputs are stacked (e.g. jnp.stack/lax.concatenate on the stack axis) while the BlockSpec's block_shape[axis] is None, Squeezed (treated as 1), or an int different from the number of stacked arrays.

Common situations: Writing a fused Pallas kernel with manual BlockSpecs and forgetting to grow the stack-axis block size when adding an input; using None block dims (unbounded) on the stacking axis; mixing squeezed and unsqueezed dims after axis juggling.

Related errors

• Block shape for {origin} (= {block_shape}) must have the sam
• 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
• Index map function {debug_info.func_src_info} for {origin} m

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