{"record":{"id":"f5d80e16b37aff0a","repo":"jax-ml/jax","slug":"multi-dot-last-dimension-of-each-array-must-match","errorCode":null,"errorMessage":"multi_dot: last dimension of each array must match first dimension of following array. Got array shapes {[a.shape for a in arrs]}","messagePattern":"multi_dot: last dimension of each array must match first dimension of following array\\. Got array shapes (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/linalg.py","lineNumber":2238,"sourceCode":"  option:\n\n  >>> jax.jit(lambda x, y, z: (x @ y) @ z).lower(x, y, z).cost_analysis()['flops']\n  600000.0\n  >>> jax.jit(lambda x, y, z: x @ (y @ z)).lower(x, y, z).cost_analysis()['flops']\n  30000.0\n  >>> jax.jit(jnp.linalg.multi_dot).lower([x, y, z]).cost_analysis()['flops']\n  30000.0\n  \"\"\"\n  arrs = list(ensure_arraylike('jnp.linalg.multi_dot', *arrays))\n  if len(arrs) < 2:\n    raise ValueError(f\"multi_dot requires at least two arrays; got len(arrays)={len(arrs)}\")\n  if not (arrs[0].ndim in (1, 2) and arrs[-1].ndim in (1, 2) and\n          all(a.ndim == 2 for a in arrs[1:-1])):\n    raise ValueError(\"multi_dot: input arrays must all be two-dimensional, except for\"\n                     \" the first and last array which may be 1 or 2 dimensional.\"\n                     f\" Got array shapes {[a.shape for a in arrs]}\")\n  if any(a.shape[-1] != b.shape[0] for a, b in zip(arrs[:-1], arrs[1:])):\n    raise ValueError(\"multi_dot: last dimension of each array must match first dimension\"\n                     f\" of following array. Got array shapes {[a.shape for a in arrs]}\")\n  einsum_axes: list[tuple[int, ...]] = [(i, i+1) for i in range(len(arrs))]\n  if arrs[0].ndim == 1:\n    einsum_axes[0] = einsum_axes[0][1:]\n  if arrs[-1].ndim == 1:\n    einsum_axes[-1] = einsum_axes[-1][:1]\n  return einsum.einsum(*itertools.chain(*zip(arrs, einsum_axes)),  # pyrefly: ignore[no-matching-overload]\n                       optimize='auto', precision=precision)\n\n\n@export\n@api.jit(static_argnames=['p'])\ndef cond(x: ArrayLike, p=None):\n  \"\"\"Compute the condition number of a matrix.\n\n  JAX implementation of :func:`numpy.linalg.cond`.\n\n  The condition number is defined as ``norm(x, p) * norm(inv(x), p)``. For ``p = 2``","sourceCodeStart":2220,"sourceCodeEnd":2256,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/linalg.py#L2220-L2256","documentation":"In a multi_dot chain, consecutive arrays must be compatible: the last dimension of each array must equal the first dimension of the next. This mirrors the contraction rule of matmul; any adjacent mismatch (a.shape[-1] != b.shape[0]) aborts with the shapes listed in the message.","triggerScenarios":"jnp.linalg.multi_dot([A, B, C]) with A shape (n, k1), B shape (k2, m) where k1 != k2; transposed matrices in the wrong orientation.","commonSituations":"Forgetting to transpose weight matrices in MLP layer chains; off-by-one dimension errors from a bad reshape earlier in the pipeline.","solutions":["Compare adjacent shapes from the error message and transpose the offending matrix (or fix its construction).","Add a preflight check: all(a.shape[-1] == b.shape[0] for a, b in zip(arrays, arrays[1:])).","Log shapes of intermediate arrays when building the chain dynamically."],"exampleFix":"// before\nout = jnp.linalg.multi_dot([x, W])  # x: (B, D), W: (H, D)\n// after\nout = jnp.linalg.multi_dot([x, W.T])  # or construct W as (D, H)","handlingStrategy":"validation","validationCode":"assert all(a.shape[-1] == b.shape[0] for a, b in zip(arrays, arrays[1:])), \\\n    [a.shape for a in arrays]\nout = jnp.linalg.multi_dot(arrays)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Preflight-check adjacent dim compatibility","Transpose weight matrices to (in, out) orientation"],"tags":["jax","numpy","linalg","matmul","shape-mismatch"],"backgroundTag":"matrix-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}