{"record":{"id":"a23f9b54f0c08275","repo":"jax-ml/jax","slug":"after-moving-axes-to-end-leading-shape-of-a-must","errorCode":null,"errorMessage":"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}","messagePattern":"After moving axes to end, leading shape of a must match shape of b\\. got a\\.shape=(.+?), b\\.shape=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/linalg.py","lineNumber":2148,"sourceCode":"    >>> a = jax.random.normal(key1, shape=(2, 2, 4))\n    >>> b = jax.random.normal(key2, shape=(2, 2))\n    >>> x = jnp.linalg.tensorsolve(a, b)\n    >>> x.shape\n    (4,)\n\n    Now show that ``x`` can be used to reconstruct ``b`` using\n    :func:`~jax.numpy.linalg.tensordot`:\n\n    >>> b_reconstructed = jnp.linalg.tensordot(a, x, axes=x.ndim)\n    >>> jnp.allclose(b, b_reconstructed)\n    Array(True, dtype=bool)\n  \"\"\"\n  a_arr, b_arr = ensure_arraylike(\"tensorsolve\", a, b)\n  if axes is not None:\n    a_arr = jnp.moveaxis(a_arr, axes, len(axes) * (a_arr.ndim - 1,))\n  out_shape = a_arr.shape[b_arr.ndim:]\n  if a_arr.shape[:b_arr.ndim] != b_arr.shape:\n    raise ValueError(\"After moving axes to end, leading shape of a must match shape of b.\"\n                     f\" got a.shape={a_arr.shape}, b.shape={b_arr.shape}\")\n  if b_arr.size != math.prod(out_shape):\n    raise ValueError(\"Input arrays must have prod(a.shape[:b.ndim]) == prod(a.shape[b.ndim:]);\"\n                     f\" got a.shape={a_arr.shape}, b.ndim={b_arr.ndim}.\")\n  a_arr = a_arr.reshape(b_arr.size, math.prod(out_shape))\n  return solve(a_arr, b_arr.ravel()).reshape(out_shape)\n\n\n@export\ndef multi_dot(arrays: Sequence[ArrayLike], *, precision: lax.PrecisionLike = None) -> Array:\n  \"\"\"Efficiently compute matrix products between a sequence of arrays.\n\n  JAX implementation of :func:`numpy.linalg.multi_dot`.\n\n  JAX internally uses the opt_einsum library to compute the most efficient\n  operation order.\n\n  Args:","sourceCodeStart":2130,"sourceCodeEnd":2166,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/linalg.py#L2130-L2166","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Reorder a's axes (or pass axes=) so a.shape[:b.ndim] == b.shape.","Fix b's shape if the right-hand side was constructed incorrectly.","Check shapes before the call: assert a.shape[:b.ndim] == b.shape."],"exampleFix":"// before\nx = jnp.linalg.tensorsolve(a, b)  # a.shape=(3, 2, 6), b.shape=(2, 3)\n// after\nx = jnp.linalg.tensorsolve(a, b, axes=(1, 0))  # moveaxis makes leading shape (2, 3) == b.shape","handlingStrategy":"validation","validationCode":"a, b = jnp.asarray(a), jnp.asarray(b)\nif a.shape[:b.ndim] != b.shape:\n    a = jnp.moveaxis(a, axes or (), range(b.ndim))\nassert a.shape[:b.ndim] == b.shape, (a.shape, b.shape)\nx = jnp.linalg.tensorsolve(a, b, axes=axes)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Verify leading shape of a matches b after any moveaxis","Print both shapes when porting NumPy examples"],"tags":["jax","numpy","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"}