why jax.scipy.sparse.linalg.cg requires the same shape for x and b

[pcasl]

when we solve the general linear problem of Ax=b, why we expect the x and b have the same shape?

It is mentioned in https://docs.jax.dev/en/latest/_autosummary/jax.scipy.sparse.linalg.cg.html as

$A$ must represent a hermitian, positive definite matrix, and must return array(s) with the same structure and shape as its argument.

And in the scipy/sparse/linalg.py line 213

  if _shapes(x0) != _shapes(b):
    raise ValueError(
        'arrays in x0 and b must have matching shapes: '
        f'{_shapes(x0)} vs {_shapes(b)}')

[pcasl]

I see. It is not a least square solver. Then which is the right least square solver in Jax? if A is an operator instead of matrix? Thanks

[patrick-kidger]

For a dense algorithm, then QR or SVD are the typical algorithms for solving least-squares. For an iterative solver (like CG), then LSMR would be a possible choice.

I think JAX only has SVD built-in by default. For more advanced use-cases / for more opinionated users, we have all of QR/SVD/LSMR over in Lineax.

[jakevdp]

I'll note that scipy's version of conjugate gradient also requires a square symmetric/hermitian matrix: https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.cg.html

[pcasl]

Thanks @jakevdp, all the sparse function in jax.scipy.sparse.linalg requires the input and output have the same shape. I am looking some function like fmincon in Matlab, and LSMR only takes A as a matrix, instead of an operator. @patrick-kidger I will check your library. Thanks!