jax-ml/jax · error · ValueError

expected search dim * 5 < matrix dim (got {k * 5}, {n})

Error message

expected search dim * 5 < matrix dim (got {k * 5}, {n})

What it means

LOBPCG in jax.experimental.sparse.linalg requires the search dimension k to satisfy k*5 < n where n is the matrix dimension. The algorithm needs the block to be small relative to the problem size for convergence guarantees and efficiency.

Source

Thrown at jax/experimental/sparse/linalg.py:256

    state = jax.lax.while_loop(cond, body, state)
    diagnostics = None
  i, X, _P, _R, _converged, theta = state

  if debug:
    assert diagnostics is not None
    return theta[0, :], X, i, diagnostics
  return theta[0, :], X, i


def _check_inputs(A, X):
  n, k = X.shape
  dt = X.dtype

  if k == 0:
    raise ValueError(f'must have search dim > 0, got {k}')

  if k * 5 >= n:
    raise ValueError(f'expected search dim * 5 < matrix dim (got {k * 5}, {n})')

  test_output = A(jnp.zeros((n, 1), dtype=X.dtype))

  if test_output.dtype != dt:
    raise ValueError(
        f'A, X must have same dtypes (were {test_output.dtype}, {dt})')

  if test_output.shape != (n, 1):
    s = test_output.shape
    raise ValueError(f'A must be ({n}, {n}) matrix A, got output {s}')


def _mm(a, b, precision=jax.lax.Precision.HIGHEST):
  return jax.lax.dot(a, b, precision=(precision, precision))

def _generate_diagnostics(prev_XPR, X, P, R, theta, converged, adj_resid):
  k = X.shape[1]
  assert X.shape == P.shape

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce the number of requested eigenpairs k so that k < n/5
2. Use dense eigendecomposition (jnp.linalg.eigh) on the todense() matrix if most eigenpairs are needed
3. Increase the problem size n if it was accidentally truncated

Example fix

// before
theta, U, _ = lobpcg_standard(A, jnp.zeros((100, 25)))  # 25*5 >= 100
// after
theta, U, _ = lobpcg_standard(A, jnp.zeros((100, 10)))  # 10*5 < 100

n, k = X.shape
assert 0 < k and k * 5 < n, f'require k*5 < n, got k={k}, n={n}'

When it happens

Trigger: Calling lobpcg_standard with a block X whose width k is too large relative to the number of rows n (k*5 >= n).

Common situations: Requesting a large fraction of eigenpairs with LOBPCG; small test matrices; if you need many eigenpairs use dense jax.numpy.linalg.eigh.

Related errors

• must have search dim > 0, got {k}
• Expected metadata dtype to be uint2, got: {meta.dtype}
• bcoo_slice: indices must have size mat.ndim={mat.ndim}
• len(strides) = {len(strides)}; expected {mat.ndim}
• bcoo_dynamic_slice: indices must have size mat.ndim={mat.ndi

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