{"record":{"id":"25c784d12e5b6f43","repo":"jax-ml/jax","slug":"stacking-only-supported-when-the-block-size-along","errorCode":null,"errorMessage":"Stacking only supported when the block size along the stack axis equals the number of inputs. Got block_dim={block_dim}, expected {n}.","messagePattern":"Stacking only supported when the block size along the stack axis equals the number of inputs\\. Got block_dim=(.+?), expected (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":1875,"sourceCode":"def _stack_pull_rule(\n    ctx: PullRuleContext,\n    block_transform: BlockIndexTransform,\n    *,\n    axis: int,\n):\n  block_shape = block_transform.block_shape\n  is_element_block = [isinstance(bd, pallas_core.Element) for bd in block_shape]\n  if any(is_element_block):\n    raise NotImplementedError(\n        'Stack with Element indexing is not yet supported.'\n    )\n  block_dim = block_shape[axis]\n  if block_dim is None or isinstance(block_dim, pallas_core.Squeezed):\n    block_dim = 1\n\n  n = len(ctx.avals_in)\n  if block_dim != n:\n    raise NotImplementedError(\n        \"Stacking only supported when the block size along the stack axis \"\n        f\"equals the number of inputs. Got block_dim={block_dim}, expected {n}.\"\n    )\n\n  new_block_shape = list(block_transform.block_shape)\n  new_block_shape.pop(axis)\n\n  def make_block_transform(child_index: int):\n    def new_block_index_transform(*idxs):\n      idx = list(block_transform.block_index_transform(*idxs))\n      idx.pop(axis)\n      return tuple(idx)\n\n    return block_transform.replace(\n        block_shape=tuple(new_block_shape),\n        block_index_transform=new_block_index_transform\n    )\n","sourceCodeStart":1857,"sourceCodeEnd":1893,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L1857-L1893","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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)","Split the stack/concatenate outside the fused kernel and pass inputs as separate block-mapped arguments","Check that the stack axis is not marked with a None (unbounded) or Squeezed block dim before fusion"],"exampleFix":"// before\nspec = BlockSpec((None, 128, 128), index_map=lambda i, j, k: (i, j, k))\nout = fuse(jnp.stack, ...)([a, b], axis=0)\n// after\nspec = BlockSpec((2, 128, 128), index_map=lambda i, j, k: (i, j, k))  # 2 == number of inputs\n","handlingStrategy":"validation","validationCode":"n = len(inputs)\nassert 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'","typeGuard":"def stack_spec_ok(block_shape, axis, inputs) -> bool:\n    bd = block_shape[axis]\n    return bd is None and len(inputs) == 1 or bd == len(inputs)","tryCatchPattern":"try:\n    out = fused_kernel(*inputs)\nexcept NotImplementedError as e:\n    if 'Stacking only supported' in str(e):\n        out = jnp.stack(inputs, axis=axis)  # unfused fallback\n    else:\n        raise","preventionTips":["Keep stack-axis block size synced with the number of stacked inputs","Avoid None/Squeezed dims on axes used for stacking","Add a unit test asserting BlockSpec matches operand count"],"tags":["jax","pallas","fusion","block-spec","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}