Internal change

PiperOrigin-RevId: 418683539

Author: tree-math authors

README.md

@@ -14,23 +14,20 @@ unlike the [implementation in JAX](https://github.com/google/jax/blob/b5aea7bc2d
 ```python
 import functools
 from jax import lax
-import tree_math
+import tree_math as tm
 import tree_math.numpy as tnp
 
-@functools.partial(tree_math.wrap, vector_argnames=['b', 'x0'])
+@functools.partial(tm.wrap, vector_argnames=['b', 'x0'])
 def cg(A, b, x0, M=lambda x: x, maxiter=5, tol=1e-5, atol=0.0):
   """jax.scipy.sparse.linalg.cg, written with tree_math."""
-  A = tree_math.unwrap(A)
-  M = tree_math.unwrap(M)
+  A = tm.unwrap(A)
+  M = tm.unwrap(M)
 
-  # tolerance handling uses the "non-legacy" behavior of scipy.sparse.linalg.cg
-  bs = b @ b
-  atol2 = tnp.maximum(tol**2 * bs, atol**2)
+  atol2 = tnp.maximum(tol**2 * (b @ b), atol**2)
 
   def cond_fun(value):
     x, r, gamma, p, k = value
-    rs = r @ r
-    return (rs > atol2) & (k < maxiter)
+    return (r @ r > atol2) & (k < maxiter)
 
   def body_fun(value):
     x, r, gamma, p, k = value

setup.py

@@ -22,6 +22,7 @@
 ]
 tests_requires = [
     'absl-py',
+    'jaxlib',
     'numpy>=1.17',
     'pytest',
 ]

tree_math/__init__.py

@@ -22,3 +22,5 @@
 )
 from tree_math._src.vector import Vector
 import tree_math.numpy
+
+__version__ = '0.1.0'

tree_math/integration_test.py

@@ -53,20 +53,15 @@ def test_cg(self):
 
     @functools.partial(tm.wrap, vector_argnames=["b", "x0"])
     def cg(A, b, x0, M=lambda x: x, maxiter=5, tol=1e-5, atol=0.0):
+      """jax.scipy.sparse.linalg.cg, written with tree_math."""
       A = tm.unwrap(A)
       M = tm.unwrap(M)
 
-      # tolerance handling uses the "non-legacy" behavior of
-      # scipy.sparse.linalg.cg
-      bs = b @ b
-      atol2 = tnp.maximum(tol**2 * bs, atol**2)
-
-      # https://en.wikipedia.org/wiki/Conjugate_gradient_method#The_preconditioned_conjugate_gradient_method
+      atol2 = tnp.maximum(tol**2 * (b @ b), atol**2)
 
       def cond_fun(value):
-        _, r, _, _, k = value
-        rs = r @ r
-        return (rs > atol2) & (k < maxiter)
+        x, r, gamma, p, k = value  # pylint: disable=unused-variable
+        return (r @ r > atol2) & (k < maxiter)
 
       def body_fun(value):
         x, r, gamma, p, k = value