jax-ml/jax · error · ValueError

tensorinv is only possible when the product of the first `in

Error message

tensorinv is only possible when the product of the first `ind` dimensions equals that of the remaining dimensions. got {arr.shape=} with {ind=}.

What it means

tensorinv treats a as a 'linear operator' mapping the last a.ndim-ind axes to the first ind axes; it reshapes a to shape (prod(shape[:ind]), prod(shape[ind:])) and inverts that matrix. The reshape only yields a square matrix — and is only invertible in this sense — when the products of the two axis-group sizes are equal.

Source

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

    - :func:`jax.numpy.linalg.tensordot`
    - :func:`jax.numpy.linalg.tensorsolve`

  Examples:
    >>> key = jax.random.key(1337)
    >>> x = jax.random.normal(key, shape=(2, 2, 4))
    >>> xinv = jnp.linalg.tensorinv(x, 2)
    >>> xinv_x = jnp.linalg.tensordot(xinv, x, axes=2)
    >>> jnp.allclose(xinv_x, jnp.eye(4), atol=1E-4)
    Array(True, dtype=bool)
  """
  arr = ensure_arraylike("tensorinv", a)
  ind = operator.index(ind)
  if ind <= 0:
    raise ValueError(f"ind must be a positive integer; got {ind=}")
  contracting_shape, batch_shape = arr.shape[:ind], arr.shape[ind:]
  flatshape = (math.prod(contracting_shape), math.prod(batch_shape))
  if flatshape[0] != flatshape[1]:
    raise ValueError("tensorinv is only possible when the product of the first"
                     " `ind` dimensions equals that of the remaining dimensions."
                     f" got {arr.shape=} with {ind=}.")
  return inv(arr.reshape(flatshape)).reshape(*batch_shape, *contracting_shape)


@export
def tensorsolve(a: ArrayLike, b: ArrayLike, axes: tuple[int, ...] | None = None) -> Array:
  """Solve the tensor equation a x = b for x.

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

  Args:
    a: input array. After reordering via ``axes`` (see below), shape must be
      ``(*b.shape, *x.shape)``.
    b: right-hand-side array.
    axes: optional tuple specifying axes of ``a`` that should be moved to the end

  Returns:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Adjust ind so the product of the first ind dims equals the product of the rest (read arr.shape from the error message).
2. Fix the tensor's construction so its input and output spaces have equal total dimension.
3. If you actually want a pseudo-inverse, use jnp.linalg.pinv on the manually reshaped matrix instead.

Example fix

// before
inv = jnp.linalg.tensorinv(jnp.zeros((2, 3, 4)), ind=2)  # 6 != 4
// after
inv = jnp.linalg.tensorinv(jnp.zeros((2, 3, 4)), ind=1)  # 2 == 3*4 -> no; use shape (2,3,2,3), ind=2

import math
if math.prod(a.shape[:ind]) != math.prod(a.shape[ind:]):
    raise ValueError(f'tensor not square-split: {a.shape=} {ind=}')
inv = jnp.linalg.tensorinv(a, ind)

When it happens

Trigger: Calling jnp.linalg.tensorinv(a, ind=k) where prod(a.shape[:k]) != prod(a.shape[k:]); e.g. shape (4, 6) with ind=1, or shape (2, 3, 4) with ind=2 (6 != 4).

Common situations: Miscounting which axes the split index covers (SciPy uses the same ind convention, so ported code with the wrong ind trips it); constructing a tensor whose legs don't represent equal-dimensional domain/codomain.

Related errors

• ind must be a positive integer; got {ind=}
• After moving axes to end, leading shape of a must match shap
• Input arrays must have prod(a.shape[:b.ndim]) == prod(a.shap
• Argument to symmetric eigendecomposition must have shape [..
• Argument to Hessenberg reduction must have shape [..., n, n]

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