jax-ml/jax · error · ValueError

ormqr with left=True expects c to have the same number of ro

Error message

ormqr with left=True expects c to have the same number of rows as the Householder matrix a. Got a shape {a_shape} and c shape {c_shape}.

What it means

jax/_src/lax/linalg.py:1527 in _ormqr_shape_rule. ormqr(a, taus, c, left=True) applies the Householder product Q (encoded in a, shape (m, n)) from the left, so c must have exactly m rows. A mismatch between c.shape[0] and a.shape[0] raises this ValueError.

Source

Thrown at jax/_src/lax/linalg.py:1527

    >>> h, taus = jnp.linalg.qr(a, mode="raw")
    >>> c = jnp.eye(3)
    >>> Q_times_c = ormqr(h.mT, taus, c)
    >>> Q_direct, _ = jnp.linalg.qr(a, mode="complete")
    >>> jnp.allclose(Q_times_c, Q_direct, atol=1e-5)
    Array(True, dtype=bool)

  See also:
    - :func:`jax.scipy.linalg.qr_multiply`: Higher-level API for computing
      Q @ C or C @ Q from a matrix ``a`` directly.
  """
  a, taus, c = core.auto_insert_reshard(a, taus, c)
  return ormqr_p.bind(a, taus, c, left=left, transpose=transpose)


def _ormqr_shape_rule(a_shape, taus_shape, c_shape, *, left, transpose):
  m = a_shape[0]
  if left and c_shape[0] != m:
    raise ValueError(
      "ormqr with left=True expects c to have the same number of rows as "
      f"the Householder matrix a. Got a shape {a_shape} and c shape {c_shape}.")
  if not left and c_shape[1] != m:
    raise ValueError(
      "ormqr with left=False expects c to have the same number of columns as "
      f"the Householder matrix a has rows. Got a shape {a_shape} and c shape {c_shape}.")
  return c_shape


@config.default_matmul_precision("highest")
def _ormqr_lowering(a, taus, c, *, left, transpose):
  # Apply Householder reflectors H_i = I - tau_i * v_i * v_i^H directly to c
  # without materializing Q. Cost: O(k * m * c_cols) if left,
  # O(k * c_rows * m) otherwise, where c has shape (..., c_rows, c_cols).
  *batch_dims, m, n = a.shape
  k = taus.shape[-1]
  is_complex = dtypes.issubdtype(a.dtype, np.complexfloating)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Reshape or pad/slice c so c.shape[0] == a.shape[0] before calling
  2. If you only need Q, use jnp.linalg.qr(a) directly which returns it
  3. Double-check which axis left=True multiplies: Q @ c requires row match

Example fix

// before
y = jax.lax.linalg.ormqr(a, taus, c, left=True)  # c.shape[0] != a.shape[0]
// after
c = jnp.pad(c, ((0, a.shape[0] - c.shape[0]),) + ((0, 0),) * (c.ndim - 1))
y = jax.lax.linalg.ormqr(a, taus, c, left=True)
Defensive patterns

Strategy: validation

Validate before calling

assert c.shape[0] == a.shape[0], (a.shape, c.shape)

Type guard

def ormqr_left_ok(a, c) -> bool:
    return c.shape[0] == a.shape[0]

Prevention

When it happens

Trigger: Calling jax.lax.linalg.ormqr with left=True where the operand c has a different row count than the reflector matrix a — e.g. applying Q to a matrix of incompatible batch or row shape, or forgetting to broadcast c correctly after qr of a transposed matrix.

Common situations: Applying Q from a QR of A to the right-hand side of a different size in custom solvers; applying Q to multiple blocks with inconsistent shapes; porting LAPACK ormqr calls with m/n arguments mixed up.

Related errors


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