jax-ml/jax · error · ValueError

The second argument to householder_product must not have mor

Error message

The second argument to householder_product must not have more rows than the minimum of the first argument's rows and columns.

What it means

jax/_src/lax/linalg.py:1429 in _householder_product_shape_rule. The taus vector (Householder scalars) must have length k <= min(m, n) of the reflector matrix a, since at most min(m, n) reflectors exist. Passing a longer taus (e.g. full-length diagonal or extra scalars) raises this ValueError.

Source

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

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(
      mlir.aval_to_ir_types(ctx.module_context, aval_out))
  op = mlir.custom_call(
      "ProductOfElementaryHouseholderReflectors",
      result_types=flat_res_types,
      operands=[a, taus],

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Trim taus to k = min(m, n): taus = taus[:min(a.shape[0], a.shape[1])]
2. Use the exact tau output returned by the factorization routine that produced a
3. When batching with padding, mask/trim tau to each matrix's reflector count

Example fix

// before
q = jax.lax.linalg.householder_product(a, taus)  # taus longer than min(m, n)
// after
k = min(a.shape[0], a.shape[1])
q = jax.lax.linalg.householder_product(a, taus[:k])

k = min(a.shape[0], a.shape[1])
taus = taus[:k]

def valid_taus(a: jax.Array, taus: jax.Array) -> bool:
    return taus.shape[-1] <= min(a.shape[-2], a.shape[-1])

When it happens

Trigger: Calling jax.lax.linalg.householder_product(a, taus) where taus.shape[0] > min(a.shape[0], a.shape[1]); e.g. taking taus from hessenberg (length n-1) but passing an n-length slice, or concatenating tau arrays across batches incorrectly.

Common situations: Manually assembling reflector data instead of using the paired outputs of geqrf/hessenberg; batching where tau vectors from different matrices were padded to the max length.

Related errors

• The first argument to householder_product must have at least
• Argument to symmetric eigendecomposition must have shape [..
• Argument to Hessenberg reduction must have shape [..., n, n]
• 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

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