{"record":{"id":"fc67e2a7c45e62f7","repo":"jax-ml/jax","slug":"a-must-be-n-n-matrix-a-got-output-s","errorCode":null,"errorMessage":"A must be ({n}, {n}) matrix A, got output {s}","messagePattern":"A must be \\((.+?), (.+?)\\) matrix A, got output (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/linalg.py","lineNumber":266,"sourceCode":"def _check_inputs(A, X):\n  n, k = X.shape\n  dt = X.dtype\n\n  if k == 0:\n    raise ValueError(f'must have search dim > 0, got {k}')\n\n  if k * 5 >= n:\n    raise ValueError(f'expected search dim * 5 < matrix dim (got {k * 5}, {n})')\n\n  test_output = A(jnp.zeros((n, 1), dtype=X.dtype))\n\n  if test_output.dtype != dt:\n    raise ValueError(\n        f'A, X must have same dtypes (were {test_output.dtype}, {dt})')\n\n  if test_output.shape != (n, 1):\n    s = test_output.shape\n    raise ValueError(f'A must be ({n}, {n}) matrix A, got output {s}')\n\n\ndef _mm(a, b, precision=jax.lax.Precision.HIGHEST):\n  return jax.lax.dot(a, b, precision=(precision, precision))\n\ndef _generate_diagnostics(prev_XPR, X, P, R, theta, converged, adj_resid):\n  k = X.shape[1]\n  assert X.shape == P.shape\n\n  diagdiag = lambda x: jnp.diag(jnp.diag(x))\n  abserr = lambda x: jnp.abs(x).sum() / (k ** 2)\n\n  XTX = _mm(X.T, X)\n  DX = diagdiag(XTX)\n  orthX = abserr(XTX - DX)\n\n  PTP = _mm(P.T, P)\n  DP = diagdiag(PTP)","sourceCodeStart":248,"sourceCodeEnd":284,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/linalg.py#L248-L284","documentation":"LOBPCG probes the operator A with an (n,1) input and requires the output to be (n,1); any other shape means A is not a square n-by-n linear operator matching X's row dimension.","triggerScenarios":"Passing a matvec that returns wrong shapes: non-square matrix, batched output, flattened vector, or operator built for a different dimension than X.","commonSituations":"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.","solutions":["Verify A is square with shape matching X.shape[0]; check the matvec returns (n,) for (n,) input or (n,1) for (n,1)","Fix the sparse matrix shape or X's dimension so they agree","Remove batching/reshape inside the callable"],"exampleFix":"// before\nA = lambda v: B @ v  # B is (m, n), m != n\nX = jnp.zeros((n, k))\n// after\nA = lambda v: M @ v  # M is (n, n)\nX = jnp.zeros((n, k))","handlingStrategy":"validation","validationCode":"n = X.shape[0]\nout = A(jnp.zeros((n, 1), dtype=X.dtype))\nassert out.shape == (n, 1), f'A must map (n,1)->(n,1), got {out.shape}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Unit-test custom operators with a zero probe of the expected shape","Ensure the sparse matrix is square and matches X's row count"],"tags":["jax","sparse","lobpcg","shape-mismatch"],"backgroundTag":"operator-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}