jax-ml/jax · error · ValueError

Input arrays must have prod(a.shape[:b.ndim]) == prod(a.shap

Error message

Input arrays must have prod(a.shape[:b.ndim]) == prod(a.shape[b.ndim:]); got a.shape={a_arr.shape}, b.ndim={b_arr.ndim}.

What it means

Even when a's leading shape matches b, tensorsolve requires the remaining (output) axes of a to have total size equal to b.size, because b is flattened and solved against the matrix reshaped to (b.size, prod(out_shape)). If prod(a.shape[b.ndim:]) != b.size the linear system has the wrong number of unknowns/equations.

Source

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

    >>> x.shape
    (4,)

    Now show that ``x`` can be used to reconstruct ``b`` using
    :func:`~jax.numpy.linalg.tensordot`:

    >>> b_reconstructed = jnp.linalg.tensordot(a, x, axes=x.ndim)
    >>> jnp.allclose(b, b_reconstructed)
    Array(True, dtype=bool)
  """
  a_arr, b_arr = ensure_arraylike("tensorsolve", a, b)
  if axes is not None:
    a_arr = jnp.moveaxis(a_arr, axes, len(axes) * (a_arr.ndim - 1,))
  out_shape = a_arr.shape[b_arr.ndim:]
  if a_arr.shape[:b_arr.ndim] != b_arr.shape:
    raise ValueError("After moving axes to end, leading shape of a must match shape of b."
                     f" got a.shape={a_arr.shape}, b.shape={b_arr.shape}")
  if b_arr.size != math.prod(out_shape):
    raise ValueError("Input arrays must have prod(a.shape[:b.ndim]) == prod(a.shape[b.ndim:]);"
                     f" got a.shape={a_arr.shape}, b.ndim={b_arr.ndim}.")
  a_arr = a_arr.reshape(b_arr.size, math.prod(out_shape))
  return solve(a_arr, b_arr.ravel()).reshape(out_shape)


@export
def multi_dot(arrays: Sequence[ArrayLike], *, precision: lax.PrecisionLike = None) -> Array:
  """Efficiently compute matrix products between a sequence of arrays.

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

  JAX internally uses the opt_einsum library to compute the most efficient
  operation order.

  Args:
    arrays: sequence of arrays. All must be two-dimensional, except the first
      and last which may be one-dimensional.
    precision: either ``None`` (default), which means the default precision for

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure prod(a.shape[b.ndim:]) == b.size by fixing a's trailing axes or b's size.
2. Verify with a quick assert before calling: assert math.prod(a.shape[b.ndim:]) == b.size.
3. For overdetermined systems use a least-squares solve on the reshaped matrix instead.

Example fix

// before
x = jnp.linalg.tensorsolve(a, b)  # (2,2,4) vs (2,)
// after
x = jnp.linalg.tensorsolve(a, b)  # with a.shape == (2, 4, 2): prod(trailing)=4... use a of shape (k, m, k) so system is square
// e.g. a = a.reshape(2, 2, 2) appropriately constructed

import math
assert math.prod(a.shape[b.ndim:]) == b.size, (a.shape, b.shape)
x = jnp.linalg.tensorsolve(a, b)

When it happens

Trigger: jnp.linalg.tensorsolve(a, b) where a.shape[:b.ndim] == b.shape but prod(a.shape[b.ndim:]) != b.size; e.g. a shape (2, 2, 4) with b shape (2,) — output space has size 4 but only 2 equations.

Common situations: Building the tensor a with the wrong number of output legs; assuming tensorsolve broadcasts or solves least-squares (it does not — it requires an exactly determined square operator per flattened system).

Related errors

• After moving axes to end, leading shape of a must match shap
• matrix_transpose requires at least 2 dimensions; got {ndim=}
• tensorinv is only possible when the product of the first `in
• multi_dot: input arrays must all be two-dimensional, except
• jnp.linalg.cond: input array must be at least 2D; got {arr.s

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