jax-ml/jax · error · ValueError

Arguments to batch_matmul must be at least 2D, got {}, {}

Error message

Arguments to batch_matmul must be at least 2D, got {}, {}

What it means

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.

Source

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

  Returns:
    An array where dimensions ``[start_dimension, stop_dimension)`` have been
    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)

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
out = lax.batch_matmul(v, M)  # v is 1-d
// after
out = lax.batch_matmul(v[None, :], M)  # or jnp.matmul(v, M)

Defensive patterns

Strategy: validation

Validate before calling

assert lhs.ndim >= 2 and rhs.ndim >= 2
out = lax.batch_matmul(lhs, rhs)

Type guard

def at_least_2d(a):
    import jax.numpy as jnp
    return a if a.ndim >= 2 else jnp.atleast_2d(a)

Prevention

Use jnp.matmul/@ for vector operands (handles 1-d promotion)
Keep operands >= 2-d inside vmapped matmul code

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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