docs(implicit): update docstrings

metaopt · Benjamin-eecs · Sep 22, 2022 · Sep 9, 2022 · Sep 9, 2022 · Sep 10, 2022
commit 84e06b29666eeaaaf2cc1c06ad19336010961d71
diff --git a/torchopt/_src/linalg.py b/torchopt/_src/linalg.py
@@ -183,13 +183,12 @@ def cg(
 
     Args:
         A: (tensor or tree of tensors or function)
-            2D array or function that calculates the linear map (matrix-vector
-            product) ``Ax`` when called like ``A(x)``. ``A`` must represent a
-            Hermitian, positive definite matrix, and must return array(s) with the
-            same structure and shape as its argument.
+            2D array or function that calculates the linear map (matrix-vector product) ``Ax`` when
+            called like ``A(x)``. ``A`` must represent a Hermitian, positive definite matrix, and
+            must return array(s) with the same structure and shape as its argument.
         b: (tensor or tree of tensors)
-            Right hand side of the linear system representing a single vector. Can be
-            stored as an array or Python container of array(s) with any shape.
+            Right hand side of the linear system representing a single vector. Can be stored as an
+            array or Python container of array(s) with any shape.
         x0: (tensor or tree of tensors, optional)
             Starting guess for the solution. Must have the same structure as ``b``.
         rtol: (float, optional, default: :const:`1e-5`)

diff --git a/torchopt/_src/linear_solve.py b/torchopt/_src/linear_solve.py
@@ -104,8 +104,8 @@ def _solve_normal_cg(
 ) -> TensorTree:
     """Solves the normal equation ``A^T A x = A^T b`` using conjugate gradient.
 
-    This can be used to solve Ax=b using conjugate gradient when A is not
-    Hermitian, positive definite.
+    This can be used to solve ``A x = b`` using conjugate gradient when ``A`` is not Hermitian,
+    positive definite.
 
     Args:
       matvec: product between ``A`` and a vector.