{"record":{"id":"47d2185879433d78","repo":"jax-ml/jax","slug":"tensorinv-is-only-possible-when-the-product-of-the","errorCode":null,"errorMessage":"tensorinv is only possible when the product of the first `ind` dimensions equals that of the remaining dimensions. got {arr.shape=} with {ind=}.","messagePattern":"tensorinv is only possible when the product of the first `ind` dimensions equals that of the remaining dimensions\\. got (.+?) with (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/linalg.py","lineNumber":2102,"sourceCode":"    - :func:`jax.numpy.linalg.tensordot`\n    - :func:`jax.numpy.linalg.tensorsolve`\n\n  Examples:\n    >>> key = jax.random.key(1337)\n    >>> x = jax.random.normal(key, shape=(2, 2, 4))\n    >>> xinv = jnp.linalg.tensorinv(x, 2)\n    >>> xinv_x = jnp.linalg.tensordot(xinv, x, axes=2)\n    >>> jnp.allclose(xinv_x, jnp.eye(4), atol=1E-4)\n    Array(True, dtype=bool)\n  \"\"\"\n  arr = ensure_arraylike(\"tensorinv\", a)\n  ind = operator.index(ind)\n  if ind <= 0:\n    raise ValueError(f\"ind must be a positive integer; got {ind=}\")\n  contracting_shape, batch_shape = arr.shape[:ind], arr.shape[ind:]\n  flatshape = (math.prod(contracting_shape), math.prod(batch_shape))\n  if flatshape[0] != flatshape[1]:\n    raise ValueError(\"tensorinv is only possible when the product of the first\"\n                     \" `ind` dimensions equals that of the remaining dimensions.\"\n                     f\" got {arr.shape=} with {ind=}.\")\n  return inv(arr.reshape(flatshape)).reshape(*batch_shape, *contracting_shape)\n\n\n@export\ndef tensorsolve(a: ArrayLike, b: ArrayLike, axes: tuple[int, ...] | None = None) -> Array:\n  \"\"\"Solve the tensor equation a x = b for x.\n\n  JAX implementation of :func:`numpy.linalg.tensorsolve`.\n\n  Args:\n    a: input array. After reordering via ``axes`` (see below), shape must be\n      ``(*b.shape, *x.shape)``.\n    b: right-hand-side array.\n    axes: optional tuple specifying axes of ``a`` that should be moved to the end\n\n  Returns:","sourceCodeStart":2084,"sourceCodeEnd":2120,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/linalg.py#L2084-L2120","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Adjust ind so the product of the first ind dims equals the product of the rest (read arr.shape from the error message).","Fix the tensor's construction so its input and output spaces have equal total dimension.","If you actually want a pseudo-inverse, use jnp.linalg.pinv on the manually reshaped matrix instead."],"exampleFix":"// before\ninv = jnp.linalg.tensorinv(jnp.zeros((2, 3, 4)), ind=2)  # 6 != 4\n// after\ninv = jnp.linalg.tensorinv(jnp.zeros((2, 3, 4)), ind=1)  # 2 == 3*4 -> no; use shape (2,3,2,3), ind=2","handlingStrategy":"validation","validationCode":"import math\nif math.prod(a.shape[:ind]) != math.prod(a.shape[ind:]):\n    raise ValueError(f'tensor not square-split: {a.shape=} {ind=}')\ninv = jnp.linalg.tensorinv(a, ind)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Check prod of axis groups before calling","Remember ind counts leading axes as the domain"],"tags":["jax","scipy","linalg","tensor","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}