jax-ml/jax · error · ValueError

Arguments to batch_matmul must have same ndim, got {}, {}

Error message

Arguments to batch_matmul must have same ndim, got {}, {}

What it means

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.

Source

Thrown at jax/_src/lax/lax.py:4001

    collapsed (raveled) into a single dimension.
  """
  lo, hi, _ = slice(start_dimension, stop_dimension).indices(len(operand.shape))
  if hi < lo:
    raise ValueError(f"Invalid dimension range passed to collapse: {operand.shape}"
                     f"[{start_dimension}:{stop_dimension}]")
  size = math.prod(operand.shape[lo:hi])
  new_shape = operand.shape[:lo] + (size,) + operand.shape[hi:]
  return reshape(operand, new_shape)


def batch_matmul(lhs: Array, rhs: Array,
                 precision: PrecisionLike = None) -> Array:
  """Batch matrix multiplication."""
  if _min(lhs.ndim, rhs.ndim) < 2:
    raise ValueError('Arguments to batch_matmul must be at least 2D, got {}, {}'
                     .format(lhs.ndim, rhs.ndim))
  if lhs.ndim != rhs.ndim:
    raise ValueError('Arguments to batch_matmul must have same ndim, got {}, {}'
                     .format(lhs.ndim, rhs.ndim))
  lhs_contract = (lhs.ndim - 1,)
  rhs_contract = (rhs.ndim - 2,)
  batch = tuple(range(lhs.ndim - 2))
  return dot_general(lhs, rhs, ((lhs_contract, rhs_contract), (batch, batch)),
                     precision=precision)


# These functions also exist in the XLA client library, but we treat them
# as non-primitive to maintain a smaller set of autodiff primitives.

def square(x: ArrayLike) -> Array:
  r"""Elementwise square: :math:`x^2`."""
  return square_p.bind(x)

def reciprocal(x: ArrayLike) -> Array:
  r"""Elementwise reciprocal: :math:`1 \over x`."""
  return integer_pow(x, -1)

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
out = lax.batch_matmul(A, Bstack)  # A: (n,m), Bstack: (b,m,k)
// after
out = lax.batch_matmul(A[None], Bstack)  # both 3-d
# or simply: out = A @ Bstack

Defensive patterns

Strategy: validation

Validate before calling

if lhs.ndim != rhs.ndim:
    lhs, rhs = jnp.broadcast_arrays(lhs[None] if lhs.ndim < rhs.ndim else lhs,
                                    rhs[None] if rhs.ndim < lhs.ndim else rhs)
out = lax.batch_matmul(lhs, rhs)

Prevention

Prefer jnp.matmul when batch broadcasting is needed
Match batch ranks explicitly with [None, ...] expansion

When it happens

Trigger: Calling lax.batch_matmul(A, B) where A.ndim != B.ndim, e.g. a (m,n) matrix times a (b,n,p) batched tensor.

Common situations: Assuming numpy-style broadcasting of batch dimensions; forgetting to add a leading batch axis to one side; mixing vmapped and non-vmapped operands.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/03f89e82ce0b33d7. Report an issue: GitHub.