{"record":{"id":"38330b8f9f188f4d","repo":"jax-ml/jax","slug":"input-arrays-must-have-prod-a-shape-b-ndim-p","errorCode":null,"errorMessage":"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}.","messagePattern":"Input arrays must have prod\\(a\\.shape\\[:b\\.ndim\\]\\) == prod\\(a\\.shape\\[b\\.ndim:\\]\\); got a\\.shape=(.+?), b\\.ndim=(.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/linalg.py","lineNumber":2151,"sourceCode":"    >>> 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:\n    arrays: sequence of arrays. All must be two-dimensional, except the first\n      and last which may be one-dimensional.\n    precision: either ``None`` (default), which means the default precision for","sourceCodeStart":2133,"sourceCodeEnd":2169,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/linalg.py#L2133-L2169","documentation":"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.","triggerScenarios":"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.","commonSituations":"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).","solutions":["Ensure prod(a.shape[b.ndim:]) == b.size by fixing a's trailing axes or b's size.","Verify with a quick assert before calling: assert math.prod(a.shape[b.ndim:]) == b.size.","For overdetermined systems use a least-squares solve on the reshaped matrix instead."],"exampleFix":"// before\nx = jnp.linalg.tensorsolve(a, b)  # (2,2,4) vs (2,)\n// after\nx = 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\n// e.g. a = a.reshape(2, 2, 2) appropriately constructed","handlingStrategy":"validation","validationCode":"import math\nassert math.prod(a.shape[b.ndim:]) == b.size, (a.shape, b.shape)\nx = jnp.linalg.tensorsolve(a, b)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Confirm the flattened system is square before solving","Don't use tensorsolve for overdetermined problems"],"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"}