jax-ml/jax · error · ValueError

multi_dot: last dimension of each array must match first dim

Error message

multi_dot: last dimension of each array must match first dimension of following array. Got array shapes {[a.shape for a in arrs]}

What it means

In a multi_dot chain, consecutive arrays must be compatible: the last dimension of each array must equal the first dimension of the next. This mirrors the contraction rule of matmul; any adjacent mismatch (a.shape[-1] != b.shape[0]) aborts with the shapes listed in the message.

Source

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

  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.

  JAX implementation of :func:`numpy.linalg.cond`.

  The condition number is defined as ``norm(x, p) * norm(inv(x), p)``. For ``p = 2``

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compare adjacent shapes from the error message and transpose the offending matrix (or fix its construction).
2. Add a preflight check: all(a.shape[-1] == b.shape[0] for a, b in zip(arrays, arrays[1:])).
3. Log shapes of intermediate arrays when building the chain dynamically.

Example fix

// before
out = jnp.linalg.multi_dot([x, W])  # x: (B, D), W: (H, D)
// after
out = jnp.linalg.multi_dot([x, W.T])  # or construct W as (D, H)

assert all(a.shape[-1] == b.shape[0] for a, b in zip(arrays, arrays[1:])), \
    [a.shape for a in arrays]
out = jnp.linalg.multi_dot(arrays)

When it happens

Trigger: jnp.linalg.multi_dot([A, B, C]) with A shape (n, k1), B shape (k2, m) where k1 != k2; transposed matrices in the wrong orientation.

Common situations: Forgetting to transpose weight matrices in MLP layer chains; off-by-one dimension errors from a bad reshape earlier in the pipeline.

Related errors

• multi_dot requires at least two arrays; got len(arrays)={len
• multi_dot: input arrays must all be two-dimensional, except
• matrix_transpose requires at least 2 dimensions; got {ndim=}
• After moving axes to end, leading shape of a must match shap
• Input arrays must have prod(a.shape[:b.ndim]) == prod(a.shap

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