{"record":{"id":"2c8af62cf3a00883","repo":"jax-ml/jax","slug":"multi-dot-input-arrays-must-all-be-two-dimensiona","errorCode":null,"errorMessage":"multi_dot: input arrays must all be two-dimensional, except for the first and last array which may be 1 or 2 dimensional. Got array shapes {[a.shape for a in arrs]}","messagePattern":"multi_dot: input arrays must all be two-dimensional, except for the first and last array which may be 1 or 2 dimensional\\. Got array shapes (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/linalg.py","lineNumber":2234,"sourceCode":"  Array(True, dtype=bool)\n\n  We can use JAX's :ref:`ahead-of-time-lowering` tools to estimate the total flops\n  of each approach, and confirm that ``multi_dot`` is choosing the more efficient\n  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.","sourceCodeStart":2216,"sourceCodeEnd":2252,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/linalg.py#L2216-L2252","documentation":"multi_dot only supports chains of 2-D matrices, with the exception that the first and last elements may be 1-D vectors (treated like NumPy matmul vector promotion). Any middle array with ndim != 2, or a first/last array with ndim not in (1, 2) — including batched 3-D+ arrays — raises this error.","triggerScenarios":"jnp.linalg.multi_dot([a, batched_b, c]) where batched_b.ndim == 3; a first array that is a scalar or 3-D batch; any 1-D array in the middle of the chain.","commonSituations":"Expecting multi_dot to vmap/broadcast over batch dimensions (it does not — use jax.vmap or stacked jnp.matmul); mixing a 1-D bias-like vector into the middle of a chain.","solutions":["Keep all middle operands 2-D; move vectors to the ends or drop them from the chain.","For batches, apply jax.vmap(jnp.linalg.multi_dot) or loop jnp.matmul over the batch.","Squeeze/reshape stray extra dims (e.g. (1, n, m) -> (n, m)) if they're accidental."],"exampleFix":"// before\nout = jnp.linalg.multi_dot([x, W1, b, W2])  # b is 1-D and in the middle\n// after\nh = jnp.matmul(x, W1) + b\nout = jnp.matmul(h, W2)","handlingStrategy":"validation","validationCode":"assert all(a.ndim == 2 for a in arrays[1:-1])\nassert arrays[0].ndim in (1, 2) and arrays[-1].ndim in (1, 2)\nout = jnp.linalg.multi_dot(arrays)","typeGuard":"def all_valid_multi_dot(arrays) -> bool:\n    return (len(arrays) >= 2 and arrays[0].ndim in (1, 2)\n            and arrays[-1].ndim in (1, 2)\n            and all(a.ndim == 2 for a in arrays[1:-1]))","tryCatchPattern":null,"preventionTips":["Keep 1-D vectors only at chain ends","Use vmap for batched chains instead of 3-D arrays"],"tags":["jax","numpy","linalg","matmul","shape-validation"],"backgroundTag":"invalid-shape-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}