{"record":{"id":"37443bc298f383c5","repo":"jax-ml/jax","slug":"concatenate-dimension-out-of-bounds-dimension","errorCode":null,"errorMessage":"concatenate dimension out of bounds: dimension {} for shapes {}.","messagePattern":"concatenate dimension out of bounds: dimension (.+?) for shapes (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7241,"sourceCode":"          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]\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):","sourceCodeStart":7223,"sourceCodeEnd":7259,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7223-L7259","documentation":"The `dimension` (axis) argument to concatenate must satisfy 0 <= dimension < operands[0].ndim. JAX validates this explicitly and raises a TypeError showing the bad dimension and all operand shapes; negative indices are not accepted at this level.","triggerScenarios":"lax.concatenate(ops, dimension=2) on 2-D arrays; passing a computed axis that equals ndim; passing a negative axis directly to jax.lax.concatenate.","commonSituations":"Hardcoding an axis then changing array rank in refactoring; computing axis from a config variable that drifted; porting NumPy code that allowed axis=-1 while calling lax directly.","solutions":["Normalize the axis first: dimension = dimension % operands[0].ndim (or use len(shape)+dimension for negatives)","Use jnp.concatenate, which supports negative axes, instead of jax.lax.concatenate","Check ndim of the first operand and clamp/validate the axis before the call"],"exampleFix":"# before\nout = jax.lax.concatenate(arrs, dimension=-1)  # raises\n# after\ndim = -1 % arrs[0].ndim\nout = jax.lax.concatenate(arrs, dimension=dim)","handlingStrategy":"validation","validationCode":"dim = dimension % operands[0].ndim  # normalize negatives\nassert 0 <= dim < operands[0].ndim","typeGuard":"def valid_axis(axis: int, ndim: int) -> bool:\n    return 0 <= axis < ndim","tryCatchPattern":null,"preventionTips":["Normalize axes with modulo before calling lax-level ops","Prefer jnp.concatenate, which handles negative axes"],"tags":["jax","concatenate","axis-out-of-bounds","shape-validation"],"backgroundTag":"axis-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}