{"record":{"id":"875716945bb0f78c","repo":"jax-ml/jax","slug":"shape-along-concat-dimension-dimension-must-be-d","errorCode":null,"errorMessage":"Shape along concat dimension {dimension} must be divisible by the block shape {block_shape[dimension]} for all children. Got shape {aval.shape}.","messagePattern":"Shape along concat dimension (.+?) must be divisible by the block shape (.+?) for all children\\. Got shape (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":1758,"sourceCode":"  is_element_block = [isinstance(bd, pallas_core.Element) for bd in block_shape]\n  if any(is_element_block):\n    raise NotImplementedError(\n        'Concatenation with Element indexing is not yet supported.'\n    )\n  block_dim = block_shape[dimension]\n  if block_dim is None:\n    block_dim = 1\n\n  if block_dim == sum(aval.shape[dimension] for aval in ctx.avals_in):\n    # Handle special case if the block contains all of the concatenated\n    # array.\n    return jax.lax.concatenate(args, dimension=dimension)\n\n  num_blocks = []\n  for aval in ctx.avals_in:\n    assert isinstance(aval, core.ShapedArray)\n    if aval.shape[dimension] % block_dim != 0:\n      raise ValueError(\n          f'Shape along concat dimension {dimension} must be divisible by the'\n          f' block shape {block_shape[dimension]} for all children. Got shape'\n          f' {aval.shape}.'\n      )\n    num_blocks.append(aval.shape[dimension] // block_dim)\n  ends = np.cumsum(num_blocks).astype(np.int32)\n  starts = np.concatenate(([0], ends[:-1])).astype(np.int32)\n\n  block_indices = ctx.get_out_block_indices()[0]\n  block_idx = block_indices[dimension]\n  valid_index = 0\n  for i in range(len(ctx.avals_in)):\n    start, end = starts[i], ends[i]\n    is_valid = (start <= block_idx) & (block_idx < end)\n    valid_index = jax.lax.select(is_valid, i, valid_index)\n  out_dtype = args[0].dtype\n  args = [a.astype(jnp.float32) if a.dtype == jnp.bfloat16 else a for a in args]\n  valid_block = jax.lax.select_n(valid_index, *args)","sourceCodeStart":1740,"sourceCodeEnd":1776,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L1740-L1776","documentation":"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.","triggerScenarios":"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.","commonSituations":"Concatenating arrays with remainder sizes after choosing block sizes that only divide the total; changing block sizes during tuning; concatenating differently-sized buffers.","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"],"exampleFix":"# before\n# children of length 10 and 6, block size 4 -> 10 % 4 != 0\nout[...] = jnp.concatenate([x, y], axis=0)\n\n# after\n# block size 2 divides both 10 and 6\nBlockSpec(block_shape=(2, ...), ...)\nout[...] = jnp.concatenate([x, y], axis=0)","handlingStrategy":"validation","validationCode":"import math\nblock_dim = block_shape[dimension]\nassert all(a.shape[dimension] % block_dim == 0 for a in arrays), \\\n    'each child extent must be divisible by the concat-axis block size'","typeGuard":null,"tryCatchPattern":"try:\n    kernel = pallas_call(fn, out_spec=BlockSpec(block_shape=(bs, ...)))(\n        jnp.concatenate, ...)  #示意\nexcept ValueError as e:\n    if 'divisible by the block shape' in str(e):\n        bs = math.gcd(*(a.shape[dimension] for a in arrays))\n        # retry with gcd block size\n    else:\n        raise","preventionTips":["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"],"tags":["jax","pallas","concatenate","block-divisibility","valueerror"],"backgroundTag":"jax-pallas-block-divisibility","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}