* touch up

jax/_src/scipy/sparse/linalg.py

@@ -225,8 +225,7 @@ def _isolve(_isolve_solve, A, b, x0=None, *, tol=1e-5, atol=0.0,
   return x, info
 
 
-def cg(A, b, x0=None, *, tol=1e-5, atol=0.0, maxiter=None, M=None,
-       assume_ipart_is_zero: bool = False):
+def cg(A, b, x0=None, *, tol=1e-5, atol=0.0, maxiter=None, M=None, assume_ipart_is_zero: bool = False):
   """Use Conjugate Gradient iteration to solve ``Ax = b``.
 
   The numerics of JAX's ``cg`` should exact match SciPy's ``cg`` (up to
@@ -275,7 +274,9 @@ def cg(A, b, x0=None, *, tol=1e-5, atol=0.0, maxiter=None, M=None,
       to reach a given error tolerance.
   assume_ipart_is_zero : bool, optional
       Whether the linear operator dtype can be assumed to have imag(A)=0.
-      Defaults to False for complex-value systems.
+      Defaults to False for complex-value systems. When True, then it can be
+      assumed that for complex operators A, A @ x and A.T @ x are equivalent,
+      leading to more efficient reverse-mode autodiff.
 
   See also
   --------
@@ -293,8 +294,7 @@ def real_valued(x):
     try:
       # Prefer to use the dtype of the operator, if available.
       if callable(A) and x0 is not None:
-        symmetric = all(map(real_valued,
-                            tree_leaves(jax.eval_shape(A, x0))))
+        symmetric = all(map(real_valued, tree_leaves(jax.eval_shape(A, x0))))
       else:
           symmetric = real_valued(A)
     except TypeError:
@@ -721,8 +721,7 @@ def _solve(A, b):
   return x, info
 
 
-def bicgstab(A, b, x0=None, *, tol=1e-5, atol=0.0, maxiter=None, M=None,
-             symmetric=False):
+def bicgstab(A, b, x0=None, *, tol=1e-5, atol=0.0, maxiter=None, M=None, symmetric=False):
   """Use Bi-Conjugate Gradient Stable iteration to solve ``Ax = b``.
 
   The numerics of JAX's ``bicgstab`` should exact match SciPy's

tests/lax_scipy_sparse_test.py

@@ -107,7 +107,7 @@ def test_cg_against_scipy(self, shape, dtype, preconditioner, assume_ipart_is_ze
     rng = jtu.rand_default(self.rng())
     A = rand_sym_pos_def(rng, shape, dtype)
     if assume_ipart_is_zero:
-      A = A - jnp.imag(A)
+      A = A.real.astype(A.dtype)
     b = rng(shape[:1], dtype)
     M = self._fetch_preconditioner(preconditioner, A, rng=rng)
 
@@ -230,7 +230,7 @@ def test_bicgstab_against_scipy(self, shape, dtype, preconditioner, symmetric):
     rng = jtu.rand_default(self.rng())
     A = rng(shape, dtype)
     if symmetric:
-      A = A @ A.T.conj()
+      A = A + A.T
     b = rng(shape[:1], dtype)
     M = self._fetch_preconditioner(preconditioner, A, rng=rng)