{"record":{"id":"03f89e82ce0b33d7","repo":"jax-ml/jax","slug":"arguments-to-batch-matmul-must-have-same-ndim-got","errorCode":null,"errorMessage":"Arguments to batch_matmul must have same ndim, got {}, {}","messagePattern":"Arguments to batch_matmul must have same ndim, got (.+?), (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":4001,"sourceCode":"    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\ndef reciprocal(x: ArrayLike) -> Array:\n  r\"\"\"Elementwise reciprocal: :math:`1 \\over x`.\"\"\"\n  return integer_pow(x, -1)","sourceCodeStart":3983,"sourceCodeEnd":4019,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L3983-L4019","documentation":"jax.lax.batch_matmul requires lhs and rhs to have identical rank. Unlike numpy broadcasting, batch dims are paired positionally via dot_general, so mismatched ndim (e.g. 2-d @ 3-d) raises this ValueError.","triggerScenarios":"Calling lax.batch_matmul(A, B) where A.ndim != B.ndim, e.g. a (m,n) matrix times a (b,n,p) batched tensor.","commonSituations":"Assuming numpy-style broadcasting of batch dimensions; forgetting to add a leading batch axis to one side; mixing vmapped and non-vmapped operands.","solutions":["Match ranks explicitly: prepend batch dims with [None, ...] or jnp.broadcast_in_dim to align batch axes","Use jnp.matmul / the @ operator, which broadcasts batch dims","If broadcasting one matrix over a batch, expand dims then rely on matmul broadcasting"],"exampleFix":"// before\nout = lax.batch_matmul(A, Bstack)  # A: (n,m), Bstack: (b,m,k)\n// after\nout = lax.batch_matmul(A[None], Bstack)  # both 3-d\n# or simply: out = A @ Bstack","handlingStrategy":"validation","validationCode":"if lhs.ndim != rhs.ndim:\n    lhs, rhs = jnp.broadcast_arrays(lhs[None] if lhs.ndim < rhs.ndim else lhs,\n                                    rhs[None] if rhs.ndim < lhs.ndim else rhs)\nout = lax.batch_matmul(lhs, rhs)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer jnp.matmul when batch broadcasting is needed","Match batch ranks explicitly with [None, ...] expansion"],"tags":["jax","lax","batch-matmul","rank-mismatch","value-error"],"backgroundTag":"rank-mismatch-in-matmul","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}