Merge branch 'main' into jpb/diagfromjac

Author: jpbrodrick89

.github/workflows/run_tests.yml

@@ -32,7 +32,7 @@ jobs:
 
       - name: Test with pytest
         run: |
-          pytest
+          python -m tests
 
       - name: Check that documentation can be built.
         run: |

lineax/_solver/cg.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from collections.abc import Callable
 from typing import Any, TypeAlias
 
@@ -267,5 +268,17 @@ def assume_full_rank(self):
 """
 
 
-def NormalCG(*args):
-    return Normal(CG(*args))
+def NormalCG(*args, **kwargs):
+    """Deprecated helper function. Use `lx.Normal(lx.CG(...))` instead.
+
+    !!! warning "Deprecated"
+        `NormalCG(...)` is deprecated in favour of `lx.Normal(lx.CG(...))`.
+        This will be removed in some future version of Lineax.
+    """
+    warnings.warn(
+        "`NormalCG(...)` is deprecated in favour of `lx.Normal(lx.CG(...))`. "
+        "This will be removed in some future version of Lineax.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return Normal(CG(*args, **kwargs))

pyproject.toml

@@ -28,7 +28,7 @@ name = "lineax"
 readme = "README.md"
 requires-python = "~=3.10"
 urls = {repository = "https://github.com/google/lineax"}
-version = "0.0.8"
+version = "0.1.0"
 
 [project.optional-dependencies]
 dev = [

tests/helpers.py

@@ -59,6 +59,8 @@ def _construct_matrix_impl(getkey, cond_cutoff, tags, size, dtype, i):
 def construct_matrix(getkey, solver, tags, num=1, *, size=3, dtype=jnp.float64):
     if isinstance(solver, lx.Normal):
         cond_cutoff = math.sqrt(1000)
+    elif isinstance(solver, lx.LSMR):
+        cond_cutoff = 10  # it's not doing super well for some reason
     else:
         cond_cutoff = 1000
     return tuple(