jax-ml/jax · error · ValueError

The first argument to householder_product must have at least

Error message

The first argument to householder_product must have at least as many rows as columns, got shape {a_shape}

What it means

jax/_src/lax/linalg.py:1424 in _householder_product_shape_rule. householder_product(a, taus) reconstructs Q from the Householder reflectors stored in a lower-trapezoidal a; it requires a's first dim >= second dim (m >= n). If the reflector matrix has more columns than rows the ValueError fires.

Source

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

      "EQ", "SIGNED")
  return [
      _replace_not_ok_with_nan(ctx, batch_dims, ok, a, ctx.avals_out[0]),
      _replace_not_ok_with_nan(ctx, batch_dims, ok, taus, ctx.avals_out[1]),
  ]


hessenberg_p = linalg_primitive(
    _hessenberg_dtype_rule, (_float | _complex,), (2,), _hessenberg_shape_rule,
    "hessenberg", multiple_results=True)
mlir.register_lowering(hessenberg_p, _hessenberg_cpu_lowering, platform="cpu")


# Householder product

def _householder_product_shape_rule(a_shape, taus_shape, **_):
  m, n = a_shape
  if m < n:
    raise ValueError(
        "The first argument to householder_product must have at least as many "
        f"rows as columns, got shape {a_shape}")
  k = taus_shape[0]
  if k > core.min_dim(m, n):
    raise ValueError(
        "The second argument to householder_product must not have more rows "
        "than the minimum of the first argument's rows and columns.")
  return a_shape


def _householder_product_lowering(ctx, a, taus):
  aval_out, = ctx.avals_out
  if not is_constant_shape(aval_out.shape):
    result_shapes = [
        mlir.eval_dynamic_shape_as_tensor(ctx, aval_out.shape)]
  else:
    result_shapes = None
  flat_res_types, _ = mlir.ir_tree_registry.flatten(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check a.shape[0] >= a.shape[1] before the call; slice the reflector block correctly (a[:, :k])
2. Re-derive from jnp.linalg.qr directly if you only need Q
3. Verify taus has k <= min(m, n) elements matching the reflector count

Example fix

// before
q = jax.lax.linalg.householder_product(a, taus)  # a: (3, 5)
// after
q = jax.lax.linalg.householder_product(a[:, :a.shape[0]], taus)  # square/tall block

assert a.shape[0] >= a.shape[1], f'need m>=n, got {a.shape}'

def valid_reflector_shape(a: jax.Array) -> bool:
    return a.ndim >= 2 and a.shape[-2] >= a.shape[-1]

When it happens

Trigger: Calling jax.lax.linalg.householder_product with a of shape (m, n) where m < n — typically a truncated/mis-sliced output of geqrf, or applying the product for a tall Q to a wide reflector block.

Common situations: Chaining qr factorization outputs into householder_product with wrong slicing; porting LAPACK orgqr call sequences where k reflectors and ldq layouts differ; off-by-one when extracting the reflector panel from QR output.

Related errors

• The second argument to householder_product must not have mor
• ormqr with left=True expects c to have the same number of ro
• ormqr with left=False expects c to have the same number of c
• Argument to symmetric eigendecomposition must have shape [..
• Argument to Hessenberg reduction must have shape [..., n, n]

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