jax-ml/jax · error · ValueError

A must be ( , ) matrix A, got output

Error message

A must be ({n}, {n}) matrix A, got output {s}

What it means

Validation in the LOBPCG input checks: applying the matrix callback A to a test vector of shape (n, 1) must return shape (n, 1). Fires when the callback produces a differently shaped output, meaning A is not an (n, n) linear operator.

Solutions

1. Verify A is square with shape matching X.shape[0]; check the matvec returns (n,) for (n,) input or (n,1) for (n,1)
2. Fix the sparse matrix shape or X's dimension so they agree
3. Remove batching/reshape inside the callable

Example fix

// before
A = lambda v: B @ v  # B is (m, n), m != n
X = jnp.zeros((n, k))
// after
A = lambda v: M @ v  # M is (n, n)
X = jnp.zeros((n, k))

n = X.shape[0]
out = A(jnp.zeros((n, 1), dtype=X.dtype))
assert out.shape == (n, 1), f'A must map (n,1)->(n,1), got {out.shape}'

When it happens

Trigger: Passing a matvec that returns wrong shapes: non-square matrix, batched output, flattened vector, or operator built for a different dimension than X.

Common situations: Using a matrix of wrong shape (m x n with m != n); a preprocessing/reshape step inside A; forgetting to transpose; operator built from a differently-sized sparse matrix.

Related errors

• A, X must have same dtypes
• Addition between sparse matrices of different shapes.
• bcoo_dynamic_slice: indices must have size mat.ndim=
• bcoo_multiply_sparse: arrays must have same number of…
• bcoo_slice: indices must have size mat.ndim=

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

Appendix: source

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

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

  diagdiag = lambda x: jnp.diag(jnp.diag(x))
  abserr = lambda x: jnp.abs(x).sum() / (k ** 2)

  XTX = _mm(X.T, X)
  DX = diagdiag(XTX)
  orthX = abserr(XTX - DX)

  PTP = _mm(P.T, P)
  DP = diagdiag(PTP)

View on GitHub (pinned to 1e1c6a8fc0)