{"record":{"id":"3a454db5f5de4a0d","repo":"jax-ml/jax","slug":"arguments-to-batch-matmul-must-be-at-least-2d-got","errorCode":null,"errorMessage":"Arguments to batch_matmul must be at least 2D, got {}, {}","messagePattern":"Arguments to batch_matmul must be at least 2D, got (.+?), (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":3998,"sourceCode":"\n  Returns:\n    An array where dimensions ``[start_dimension, stop_dimension)`` have been\n    collapsed (raveled) into a single dimension.\n  \"\"\"\n  lo, hi, _ = slice(start_dimension, stop_dimension).indices(len(operand.shape))\n  if hi < lo:\n    raise ValueError(f\"Invalid dimension range passed to collapse: {operand.shape}\"\n                     f\"[{start_dimension}:{stop_dimension}]\")\n  size = math.prod(operand.shape[lo:hi])\n  new_shape = operand.shape[:lo] + (size,) + operand.shape[hi:]\n  return reshape(operand, new_shape)\n\n\ndef batch_matmul(lhs: Array, rhs: Array,\n                 precision: PrecisionLike = None) -> Array:\n  \"\"\"Batch matrix multiplication.\"\"\"\n  if _min(lhs.ndim, rhs.ndim) < 2:\n    raise ValueError('Arguments to batch_matmul must be at least 2D, got {}, {}'\n                     .format(lhs.ndim, rhs.ndim))\n  if lhs.ndim != rhs.ndim:\n    raise ValueError('Arguments to batch_matmul must have same ndim, got {}, {}'\n                     .format(lhs.ndim, rhs.ndim))\n  lhs_contract = (lhs.ndim - 1,)\n  rhs_contract = (rhs.ndim - 2,)\n  batch = tuple(range(lhs.ndim - 2))\n  return dot_general(lhs, rhs, ((lhs_contract, rhs_contract), (batch, batch)),\n                     precision=precision)\n\n\n# These functions also exist in the XLA client library, but we treat them\n# as non-primitive to maintain a smaller set of autodiff primitives.\n\ndef square(x: ArrayLike) -> Array:\n  r\"\"\"Elementwise square: :math:`x^2`.\"\"\"\n  return square_p.bind(x)\n","sourceCodeStart":3980,"sourceCodeEnd":4016,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L3980-L4016","documentation":"jax.lax.batch_matmul performs batched matrix multiplication and requires both operands to have ndim >= 2 (at least one matrix dimension plus batch dims). If either operand is 0-d or 1-d it raises this ValueError, unlike numpy matmul which promotes 1-d inputs.","triggerScenarios":"Calling lax.batch_matmul(vec, mat) with a 1-d vector, or scalars (0-d) as either argument; also under vmap when the batched operand ends up 1-d.","commonSituations":"Reusing code written for numpy @ / jnp.matmul with vector operands; vmap over the wrong axis of a matrix multiply collapsing a dimension; forgetting to add a batch dimension before batching a matmul.","solutions":["Promote vectors: use x[:, None] / x[None, :] so both operands are >=2-d, or use jnp.matmul which handles 1-d promotion","If you meant a dot product, use lax.dot / jnp.dot instead","Check ndim of inputs (and inside vmapped functions) before calling"],"exampleFix":"// before\nout = lax.batch_matmul(v, M)  # v is 1-d\n// after\nout = lax.batch_matmul(v[None, :], M)  # or jnp.matmul(v, M)","handlingStrategy":"validation","validationCode":"assert lhs.ndim >= 2 and rhs.ndim >= 2\nout = lax.batch_matmul(lhs, rhs)","typeGuard":"def at_least_2d(a):\n    import jax.numpy as jnp\n    return a if a.ndim >= 2 else jnp.atleast_2d(a)","tryCatchPattern":null,"preventionTips":["Use jnp.matmul/@ for vector operands (handles 1-d promotion)","Keep operands >= 2-d inside vmapped matmul code"],"tags":["jax","lax","batch-matmul","dimensionality","value-error"],"backgroundTag":"insufficient-array-rank","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}