jax-ml/jax · error · ValueError

After moving axes to end, leading shape of a must match shap

Error message

After moving axes to end, leading shape of a must match shape of b. got a.shape={a_arr.shape}, b.shape={b_arr.shape}

What it means

jnp.linalg.tensorsolve(a, b) solves a x = b for x when a is a tensor viewed as a matrix acting on b's space. After optionally moving axes, the leading b.ndim dimensions of a must exactly equal b's shape (a's 'input' legs must match b's indices). If they don't, the equation is ill-formed and this ValueError is raised.

Source

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

    >>> a = jax.random.normal(key1, shape=(2, 2, 4))
    >>> b = jax.random.normal(key2, shape=(2, 2))
    >>> x = jnp.linalg.tensorsolve(a, b)
    >>> 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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reorder a's axes (or pass axes=) so a.shape[:b.ndim] == b.shape.
2. Fix b's shape if the right-hand side was constructed incorrectly.
3. Check shapes before the call: assert a.shape[:b.ndim] == b.shape.

Example fix

// before
x = jnp.linalg.tensorsolve(a, b)  # a.shape=(3, 2, 6), b.shape=(2, 3)
// after
x = jnp.linalg.tensorsolve(a, b, axes=(1, 0))  # moveaxis makes leading shape (2, 3) == b.shape

a, b = jnp.asarray(a), jnp.asarray(b)
if a.shape[:b.ndim] != b.shape:
    a = jnp.moveaxis(a, axes or (), range(b.ndim))
assert a.shape[:b.ndim] == b.shape, (a.shape, b.shape)
x = jnp.linalg.tensorsolve(a, b, axes=axes)

When it happens

Trigger: jnp.linalg.tensorsolve(a, b) where a.shape[:b.ndim] != b.shape, e.g. a of shape (2, 3, 6) with b of shape (3, 3); also hit when the axes= argument moved the wrong axes so the leading shape no longer matches b.

Common situations: Porting np.linalg.tensorsolve examples with mismatched leg ordering; forgetting that with axes specified the check happens AFTER moveaxis, so the original leading shape is what must match after reordering.

Related errors

• Input arrays must have prod(a.shape[:b.ndim]) == prod(a.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/a23f9b54f0c08275. Report an issue: GitHub.