jax-ml/jax · error · TypeError

All objects to concatenate must be arrays, got {}.

Error message

All objects to concatenate must be arrays, got {}.

What it means

JAX's concatenate shape rule requires every operand to be a ShapedArray; passing a tracer, a Python scalar/list, None, or an opaque object triggers this TypeError naming the offending type. It guards the abstract-evaluation stage before lowering to XLA.

Source

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

          select(bitwise_and(gt(min, operand), lt(min, max)),
                 g, _zeros(operand)),
          lambda g, min, operand, max:
          select(bitwise_and(gt(operand, min), lt(operand, max)),
                 g, _zeros(operand)),
          lambda g, min, operand, max:
          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:]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert non-array operands with jnp.asarray before concatenating
2. Inspect the reported type to find which element of the operand list is not a JAX array
3. Flatten pytrees with jax.tree_util and filter/convert leaves explicitly

Example fix

// before
out = jnp.concatenate([x, maybe_none], axis=0)
// after
arrays = [jnp.asarray(a) for a in [x, maybe_none] if a is not None]
out = jnp.concatenate(arrays, axis=0)

arrays = [jnp.asarray(a) for a in operands if a is not None]

import jax.numpy as jnp

def all_arrays(xs) -> bool:
    return all(isinstance(x, jnp.ndarray) or hasattr(x, '__jax_array__') for x in xs)

try:
    out = jnp.concatenate(arrays, axis=0)
except TypeError as e:
    if 'must be arrays' in str(e):
        arrays = [jnp.asarray(a) for a in arrays]
        out = jnp.concatenate(arrays, axis=0)
    else:
        raise

When it happens

Trigger: jnp.concatenate([x, None]), mixing a Python list into the operand list, passing a pytree node or a dict value instead of a traced array.

Common situations: Unpacking tuples/pytrees where one element isn't an array; passing model config objects or None defaults into a concat call; list-of-lists data not converted via jnp.asarray.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• operand to reduce_window_sum must have a number dtype, got {
• mulhi requires integer inputs, got {dtype}
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}

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