jax-ml/jax · error · ValueError

multi_dot: input arrays must all be two-dimensional, except

Error message

multi_dot: input arrays must all be two-dimensional, except for the first and last array which may be 1 or 2 dimensional. Got array shapes {[a.shape for a in arrs]}

What it means

multi_dot only supports chains of 2-D matrices, with the exception that the first and last elements may be 1-D vectors (treated like NumPy matmul vector promotion). Any middle array with ndim != 2, or a first/last array with ndim not in (1, 2) — including batched 3-D+ arrays — raises this error.

Source

Thrown at jax/_src/numpy/linalg.py:2234

  Array(True, dtype=bool)

  We can use JAX's :ref:`ahead-of-time-lowering` tools to estimate the total flops
  of each approach, and confirm that ``multi_dot`` is choosing the more efficient
  option:

  >>> jax.jit(lambda x, y, z: (x @ y) @ z).lower(x, y, z).cost_analysis()['flops']
  600000.0
  >>> jax.jit(lambda x, y, z: x @ (y @ z)).lower(x, y, z).cost_analysis()['flops']
  30000.0
  >>> jax.jit(jnp.linalg.multi_dot).lower([x, y, z]).cost_analysis()['flops']
  30000.0
  """
  arrs = list(ensure_arraylike('jnp.linalg.multi_dot', *arrays))
  if len(arrs) < 2:
    raise ValueError(f"multi_dot requires at least two arrays; got len(arrays)={len(arrs)}")
  if not (arrs[0].ndim in (1, 2) and arrs[-1].ndim in (1, 2) and
          all(a.ndim == 2 for a in arrs[1:-1])):
    raise ValueError("multi_dot: input arrays must all be two-dimensional, except for"
                     " the first and last array which may be 1 or 2 dimensional."
                     f" Got array shapes {[a.shape for a in arrs]}")
  if any(a.shape[-1] != b.shape[0] for a, b in zip(arrs[:-1], arrs[1:])):
    raise ValueError("multi_dot: last dimension of each array must match first dimension"
                     f" of following array. Got array shapes {[a.shape for a in arrs]}")
  einsum_axes: list[tuple[int, ...]] = [(i, i+1) for i in range(len(arrs))]
  if arrs[0].ndim == 1:
    einsum_axes[0] = einsum_axes[0][1:]
  if arrs[-1].ndim == 1:
    einsum_axes[-1] = einsum_axes[-1][:1]
  return einsum.einsum(*itertools.chain(*zip(arrs, einsum_axes)),  # pyrefly: ignore[no-matching-overload]
                       optimize='auto', precision=precision)


@export
@api.jit(static_argnames=['p'])
def cond(x: ArrayLike, p=None):
  """Compute the condition number of a matrix.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Keep all middle operands 2-D; move vectors to the ends or drop them from the chain.
For batches, apply jax.vmap(jnp.linalg.multi_dot) or loop jnp.matmul over the batch.
Squeeze/reshape stray extra dims (e.g. (1, n, m) -> (n, m)) if they're accidental.

Example fix

// before
out = jnp.linalg.multi_dot([x, W1, b, W2])  # b is 1-D and in the middle
// after
h = jnp.matmul(x, W1) + b
out = jnp.matmul(h, W2)

Defensive patterns

Strategy: validation

Validate before calling

assert all(a.ndim == 2 for a in arrays[1:-1])
assert arrays[0].ndim in (1, 2) and arrays[-1].ndim in (1, 2)
out = jnp.linalg.multi_dot(arrays)

Type guard

def all_valid_multi_dot(arrays) -> bool:
    return (len(arrays) >= 2 and arrays[0].ndim in (1, 2)
            and arrays[-1].ndim in (1, 2)
            and all(a.ndim == 2 for a in arrays[1:-1]))

Prevention

Keep 1-D vectors only at chain ends
Use vmap for batched chains instead of 3-D arrays

When it happens

Trigger: jnp.linalg.multi_dot([a, batched_b, c]) where batched_b.ndim == 3; a first array that is a scalar or 3-D batch; any 1-D array in the middle of the chain.

Common situations: Expecting multi_dot to vmap/broadcast over batch dimensions (it does not — use jax.vmap or stacked jnp.matmul); mixing a 1-D bias-like vector into the middle of a chain.

Related errors

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