jax-ml/jax · error · ValueError

must have search dim > 0, got {k}

Error message

must have search dim > 0, got {k}

What it means

Input validation for sparse LOBPCG (locally optimal block preconditioned conjugate gradient): the search block X must have a positive number of columns k (eigenvectors sought). A zero-width block cannot be used to start the iteration.

Source

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

    state, diagnostics = jax.lax.scan(
        lambda state, _: body(state), state, xs=None, length=m)
  else:
    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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Ensure the number of requested eigenpairs k >= 1
  2. Check upstream computation that produced X's width before calling lobpcg_standard

Example fix

// before
X = jnp.zeros((n, 0))
theta, U, iters = lobpcg_standard(A, X)
// after
X = jnp.zeros((n, 5))
theta, U, iters = lobpcg_standard(A, X)
Defensive patterns

Strategy: validation

Validate before calling

n, k = X.shape
if k < 1:
    raise ValueError(f'need at least one eigenpair, got k={k}')

Prevention

When it happens

Trigger: Calling jax.experimental.sparse.linalg.lobpcg_standard with X of shape (n, 0), e.g. because the number of requested eigenpairs was computed as 0.

Common situations: Passing k=0 from a config; deriving k from data (e.g. k = rank estimate that evaluates to 0) inside a hyperparameter sweep.

Related errors


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