jax-ml/jax · error · ValueError

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

Error message

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

What it means

jax/_src/lax/linalg.py:1531 in _ormqr_shape_rule. With left=False, ormqr applies Q from the right (c @ Q), so c's column count (c.shape[1]) must equal the row count of the reflector matrix a (a.shape[0]). A mismatch raises this ValueError.

Source

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

    >>> 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)

  # Householder vectors: lower triangle of a with unit diagonal.
  eye = lax._eye(a.dtype, (m, k))
  if batch_dims:
    eye = lax.broadcast(eye, tuple(batch_dims))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Adjust c so c.shape[1] == a.shape[0] (transpose, pad, or slice appropriately)
  2. Consider left=True with c.T then transpose the result if that matches your math
  3. Validate shapes with an assert before the call during development

Example fix

// before
y = jax.lax.linalg.ormqr(a, taus, c, left=False)  # c.shape[1] != a.shape[0]
// after
assert c.shape[1] == a.shape[0], (c.shape, a.shape)
y = jax.lax.linalg.ormqr(a, taus, c, left=False)
Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Prevention

When it happens

Trigger: Calling jax.lax.linalg.ormqr(a, taus, c, left=False) with c.shape[1] != a.shape[0]; e.g. applying the transpose-Q to a wide RHS whose width does not match the QR'd matrix's rows.

Common situations: Solving least-squares normal equations manually; using Q^T on both sides of a rectangular problem with inconsistent shapes; ported LAPACK code where 'left'/'trans' flags were flipped.

Related errors


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