{"record":{"id":"132ef7839ed646d0","repo":"jax-ml/jax","slug":"cannot-concatenate-arrays-with-different-numbers-o","errorCode":null,"errorMessage":"Cannot concatenate arrays with different numbers of dimensions: got {}.","messagePattern":"Cannot concatenate arrays with different numbers of dimensions: got (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7238,"sourceCode":"          select(bitwise_and(gt(operand, min), lt(operand, max)),\n                 g, _zeros(operand)),\n          lambda g, min, operand, max:\n          select(lt(max, operand), g, _zeros(operand)))\nbatching.primitive_batchers[clamp_p] = _clamp_batch_rule\nmlir.register_lowering(clamp_p, partial(_nary_lower_hlo, hlo.clamp))\n\ndef _concatenate_shape_rule(*operands, **kwargs):\n  dimension = kwargs.pop('dimension')\n  if not operands:\n    msg = \"concatenate expects at least one operand, got 0.\"\n    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]","sourceCodeStart":7220,"sourceCodeEnd":7256,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7220-L7256","documentation":"Concatenation only joins arrays along one axis, so all operands must have the same number of dimensions. JAX checks that the set of ndims has size 1 and otherwise raises this TypeError listing every shape, since NumPy-style broadcasting is not applied by lax.concatenate.","triggerScenarios":"jnp.concatenate([jnp.zeros((3,)), jnp.zeros((2,3))], axis=0) — mixing rank-1 and rank-2 arrays; concatenating scalars with vectors.","commonSituations":"Appending a scalar or 1-D row to a 2-D batch without reshaping; mixed data pipelines where some tensors went through squeeze/reshape.","solutions":["Add or remove axes so all arrays match rank: use x[None, :] / jnp.expand_dims or jnp.atleast_2d","Use jnp.stack instead if you want to add a new axis","Verify intermediate shapes with prints or jax.debug.print before the concat"],"exampleFix":"# before\nout = jnp.concatenate([batch, row], axis=0)  # row has shape (n,)\n# after\nout = jnp.concatenate([batch, row[None, :]], axis=0)","handlingStrategy":"validation","validationCode":"if len({a.ndim for a in arrays}) != 1:\n    arrays = [jnp.atleast_2d(a) for a in arrays]\nout = jnp.concatenate(arrays, axis=0)","typeGuard":"def same_rank(xs) -> bool:\n    nd = xs[0].ndim if xs else None\n    return all(x.ndim == nd for x in xs)","tryCatchPattern":null,"preventionTips":["Assert uniform rank before concat in preprocessing helpers","Prefer jnp.atleast_ndim when mixing scalars/vectors into batches"],"tags":["jax","concatenate","rank-mismatch","shape-validation"],"backgroundTag":"ndim-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}