{"record":{"id":"f05f080c64846272","repo":"jax-ml/jax","slug":"cannot-concatenate-arrays-with-shapes-that-differ","errorCode":null,"errorMessage":"Cannot concatenate arrays with shapes that differ in dimensions other than the one being concatenated: concatenating along dimension {} for shapes {}.","messagePattern":"Cannot concatenate arrays with shapes that differ in dimensions other than the one being concatenated: concatenating along dimension (.+?) for shapes (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7249,"sourceCode":"    raise TypeError(msg)\n  if not all(isinstance(operand, ShapedArray) for operand in operands):\n    msg = \"All objects to concatenate must be arrays, got {}.\"\n    op = next(op for op in operands if not isinstance(op, ShapedArray))\n    raise TypeError(msg.format(type(op)))\n  if len({operand.ndim for operand in operands}) != 1:\n    msg = \"Cannot concatenate arrays with different numbers of dimensions: got {}.\"\n    raise TypeError(msg.format(\", \".join(str(o.shape) for o in operands)))\n  if not 0 <= dimension < operands[0].ndim:\n    msg = \"concatenate dimension out of bounds: dimension {} for shapes {}.\"\n    raise TypeError(msg.format(dimension, \", \".join([str(o.shape) for o in operands])))\n  shapes = [operand.shape[:dimension] + operand.shape[dimension+1:]\n            for operand in operands]\n  if shapes[:-1] != shapes[1:]:\n    msg = (\"Cannot concatenate arrays with shapes that differ in dimensions \"\n           \"other than the one being concatenated: concatenating along \"\n           \"dimension {} for shapes {}.\")\n    shapes = [operand.shape for operand in operands]\n    raise TypeError(msg.format(dimension, \", \".join(map(str, shapes))))\n\n  concat_size = sum(o.shape[dimension] for o in operands)\n  ex_shape = operands[0].shape\n  return ex_shape[:dimension] + (concat_size,) + ex_shape[dimension+1:]\n\ndef _concatenate_sharding_rule(*operands, **kwargs):\n  non_empty_s = [o.sharding for o in operands if not o.sharding.mesh.empty]\n  if not non_empty_s:\n    return core.get_cur_mesh_sharding()\n  if not all(s == non_empty_s[0] for s in non_empty_s):\n    ss = \", \".join(str(o.sharding) for o in operands)\n    raise core.ShardingTypeError(\n        f\"All operands should have the same sharding. Got shardings {ss}\")\n  return non_empty_s[0]\n\ndef _concatenate_reduced_rule(*operands, **kwargs):\n  reduced_specs = {r for o in operands if (r := getr(o))}\n  if len(reduced_specs) > 1:","sourceCodeStart":7231,"sourceCodeEnd":7267,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7231-L7267","documentation":"For concatenation along `dimension`, every non-concatenated axis must match exactly across operands. JAX computes each shape with the concat axis removed and raises this TypeError if they differ — lax.concatenate never broadcasts.","triggerScenarios":"jnp.concatenate([jnp.zeros((2,3)), jnp.zeros((2,4))], axis=0) — second axis 3 vs 4; feature dimensions differing across batches.","commonSituations":"Concatenating sequences of embeddings/timesteps where a feature dim was mistyped; ragged per-sample data naively converted to arrays; unit drift after reshapes or transposes upstream.","solutions":["Print/assert all shapes except the concat axis match before concatenating","Fix the upstream producer so the non-concat dimensions agree (correct feature size, transpose, or reshape)","For ragged data, pad to a common size or store as a list instead of concatenating"],"exampleFix":"# before\nout = jnp.concatenate([a, b], axis=0)  # a:(2,3), b:(2,4)\n# after\nassert a.shape[1:] == b.shape[1:], (a.shape, b.shape)\nout = jnp.concatenate([a, b], axis=0)","handlingStrategy":"validation","validationCode":"ref = arrays[0].shape[:axis] + arrays[0].shape[axis+1:]\nassert all(a.shape[:axis] + a.shape[axis+1:] == ref for a in arrays), \\\n    [a.shape for a in arrays]\nout = jnp.concatenate(arrays, axis=axis)","typeGuard":"def concat_compatible(xs, axis) -> bool:\n    return len({x.shape[:axis] + x.shape[axis+1:] for x in xs}) == 1","tryCatchPattern":null,"preventionTips":["Add shape assertions in test suites for all concat sites","Log shapes when assembling batches from heterogeneous sources"],"tags":["jax","concatenate","shape-mismatch","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}