jax-ml/jax · error · TypeError

concatenate dimension out of bounds: dimension {} for shapes

Error message

concatenate dimension out of bounds: dimension {} for shapes {}.

What it means

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.

Source

Thrown at jax/_src/lax/lax.py:7241

          select(lt(max, operand), g, _zeros(operand)))
batching.primitive_batchers[clamp_p] = _clamp_batch_rule
mlir.register_lowering(clamp_p, partial(_nary_lower_hlo, hlo.clamp))

def _concatenate_shape_rule(*operands, **kwargs):
  dimension = kwargs.pop('dimension')
  if not operands:
    msg = "concatenate expects at least one operand, got 0."
    raise TypeError(msg)
  if not all(isinstance(operand, ShapedArray) for operand in operands):
    msg = "All objects to concatenate must be arrays, got {}."
    op = next(op for op in operands if not isinstance(op, ShapedArray))
    raise TypeError(msg.format(type(op)))
  if len({operand.ndim for operand in operands}) != 1:
    msg = "Cannot concatenate arrays with different numbers of dimensions: got {}."
    raise TypeError(msg.format(", ".join(str(o.shape) for o in operands)))
  if not 0 <= dimension < operands[0].ndim:
    msg = "concatenate dimension out of bounds: dimension {} for shapes {}."
    raise TypeError(msg.format(dimension, ", ".join([str(o.shape) for o in operands])))
  shapes = [operand.shape[:dimension] + operand.shape[dimension+1:]
            for operand in operands]
  if shapes[:-1] != shapes[1:]:
    msg = ("Cannot concatenate arrays with shapes that differ in dimensions "
           "other than the one being concatenated: concatenating along "
           "dimension {} for shapes {}.")
    shapes = [operand.shape for operand in operands]
    raise TypeError(msg.format(dimension, ", ".join(map(str, shapes))))

  concat_size = sum(o.shape[dimension] for o in operands)
  ex_shape = operands[0].shape
  return ex_shape[:dimension] + (concat_size,) + ex_shape[dimension+1:]

def _concatenate_sharding_rule(*operands, **kwargs):
  non_empty_s = [o.sharding for o in operands if not o.sharding.mesh.empty]
  if not non_empty_s:
    return core.get_cur_mesh_sharding()
  if not all(s == non_empty_s[0] for s in non_empty_s):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Normalize the axis first: dimension = dimension % operands[0].ndim (or use len(shape)+dimension for negatives)
2. Use jnp.concatenate, which supports negative axes, instead of jax.lax.concatenate
3. Check ndim of the first operand and clamp/validate the axis before the call

Example fix

# before
out = jax.lax.concatenate(arrs, dimension=-1)  # raises
# after
dim = -1 % arrs[0].ndim
out = jax.lax.concatenate(arrs, dimension=dim)

dim = dimension % operands[0].ndim  # normalize negatives
assert 0 <= dim < operands[0].ndim

def valid_axis(axis: int, ndim: int) -> bool:
    return 0 <= axis < ndim

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• concatenate expects at least one operand, got 0.
• Cannot concatenate arrays with different numbers of dimensio
• Cannot concatenate arrays with shapes that differ in dimensi
• scan got `length` argument of {} which disagrees with leadin
• axis {} is out of bounds for array of shape {}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/37443bc298f383c5. Report an issue: GitHub.