jax-ml/jax · error · ValueError

multi_dot requires at least two arrays; got len(arrays)={len

Error message

multi_dot requires at least two arrays; got len(arrays)={len(arrs)}

What it means

jnp.linalg.multi_dot(arrays) chains matrix products with an optimized parenthesization (like np.linalg.multi_dot / scipy.linalg.blas.dgemm_seq). Chaining requires at least two arrays; passing a single array or an empty list fails this check.

Source

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

  >>> result3 = jnp.linalg.multi_dot([x, y, z])
  >>> jnp.allclose(result1, result3, atol=1E-4)
  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Handle the degenerate cases yourself: return the single array (or identity) when len < 2.
  2. Guard the call: if len(arrays) < 2 use plain matmul or return the element.
  3. Fix the list construction so it always has >= 2 arrays.

Example fix

// before
out = jnp.linalg.multi_dot(mats)  # mats may be [a]
// after
out = mats[0] if len(mats) == 1 else jnp.linalg.multi_dot(mats)
Defensive patterns

Strategy: type-guard

Validate before calling

if len(arrays) < 2:
    out = arrays[0] if arrays else None
else:
    out = jnp.linalg.multi_dot(arrays)

Type guard

def multi_dot_ok(arrays) -> bool:
    return len(arrays) >= 2

Prevention

When it happens

Trigger: jnp.linalg.multi_dot([a]) or jnp.linalg.multi_dot([]); also multi_dot(*mats) where mats happens to contain one element, or unrolling a loop that collapses to a single array.

Common situations: Dynamic product chains built from lists whose length varies and can degenerate to 0 or 1; refactoring chained @ expressions into multi_dot without preserving at least two operands.

Related errors


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