Remove deprecated jax apis (#113)

`jax.experimental.host_callback` was recently removed in [jax
0.4.35](https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-35-oct-22-2024).

> - `jax.experimental.host_callback` has been deprecated since March
2024, with JAX version 0.4.26. Now we removed it. See
jax-ml/jax#20385 for a discussion of
alternatives.

This PR replaces `jax.experimental.host_callback` with
`jax.debug.callback` and `jax.experimental.io_callback`. In addition, it
implements the batching change in `jax.numpy.linalg.solve()` since [jax
0.5.0](https://jax.readthedocs.io/en/latest/changelog.html#jax-0-5-0-jan-17-2025).

> -
[`jax.numpy.linalg.solve()`](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.solve.html#jax.numpy.linalg.solve)
no longer supports batched 1D arguments on the right hand side. To
recover the previous behavior in these cases, use `solve(a, b[...,
None]).squeeze(-1)`.

Lastly, jax's prng changed in 0.5.0 so the generated random numbers will
be different.

> Enable jax_threefry_partitionable by default (see [the update
note](jax-ml/jax#18480)).

This is unlucky on the power method test, so I've bumped the relative
tolerance slightly.

(Very untested, but at least it passes the tests. It seems tqdm has a
bit of a footgun where if a jax array is passed to info, then it'll
freeze. Perhaps this should be documented in the `kwargs` argument.)

---------

Co-authored-by: Andres Potapczynski <[email protected]>

Author: stephen-huan

README.md

@@ -22,7 +22,7 @@
 [![Downloads](https://static.pepy.tech/badge/cola-ml)](https://pepy.tech/project/cola-ml)
 <!-- [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wilson-labs/cola/blob/master/docs/notebooks/colabs/all.ipynb) -->
 
-CoLA is a framework for scalable linear algebra in machine learning and beyond, providing: 
+CoLA is a framework for scalable linear algebra in machine learning and beyond, providing:
 
 (1) Fast hardware-sensitive (GPU accelerated) iterative algorithms for general matrix operations; <br>
 (2) Algorithms that can exploit matrix structure for efficiency; <br>

cola/linalg/inverse/gmres.py

@@ -115,7 +115,7 @@ def gmres_fwd(A, rhs, x0, max_iters, tol, P, use_householder, use_triangular, pb
         zero_thresh = 10 * tol * overall_max[:, None]
         padding = xnp.where(largest_vals < zero_thresh, xnp.ones_like(largest_vals), xnp.zeros_like(largest_vals))
         added_diag = xnp.vmap(xnp.diag)(padding)
-        y = xnp.solve(HT @ H + added_diag, HT[..., 0]) * beta[:, None]
+        y = xnp.solve(HT @ H + added_diag, HT[..., 0, None]).squeeze(-1) * beta[:, None]
         zeros = xnp.zeros_like(y)
         y = xnp.where(largest_vals < zero_thresh, zeros, y)
         pred = xnp.permute(Q @ y[..., None], axes=[1, 0, 2])[:, :, 0]

cola/linalg/tbd/nullspace.py

@@ -106,7 +106,7 @@ def cond_fn(state):
     if final_error > 5 * tol:
         logging.warning(f"Normalized basis has too high error {final_error:.2e} for tol {tol:.2e}")
     scutoff = (S[rank] if r > rank else 0)
-    text = f"Singular value gap too small: {S[rank-1]:.2e}"
+    text = f"Singular value gap too small: {S[rank - 1]:.2e}"
     text += "above cutoff {scutoff:.2e} below cutoff. Final L, earlier {S[rank-5:rank]}"
     assert rank == 0 or scutoff < S[rank - 1] / 100, text
 

cola/linalg/unary/unary.py

@@ -81,7 +81,7 @@ def _matmat(self, V):  # (n,bs)
         norms = self.xnp.norm(V, axis=0)
 
         e0 = self.xnp.canonical(0, (P.shape[1], V.shape[-1]), dtype=P.dtype, device=self.device)
-        Pinv0 = self.xnp.solve(P, e0.T)  # (bs, m, m) vs (bs, m)
+        Pinv0 = self.xnp.solve(P, e0.T[..., None]).squeeze(-1)  # (bs, m, m) vs (bs, m)
         out = Pinv0 * norms[:, None]  # (bs, m)
         Q = self.xnp.cast(Q, dtype=P.dtype)  # (bs, n, m)
         # (bs,n,m) @ (bs,m,m) @ (bs, m) -> (bs, n)

cola/utils/jax_tqdm.py

@@ -6,7 +6,8 @@
 
 import jax
 import numpy as np
-from jax.experimental import host_callback
+from jax.debug import callback
+from jax.experimental import io_callback
 from tqdm.auto import tqdm
 
 
@@ -94,45 +95,46 @@ def build_tqdm(n: int, message: typing.Optional[str] = None) -> typing.Tuple[typ
         print_rate = 1
     remainder = n % print_rate
 
-    def _define_tqdm(arg, transform):
+    def _define_tqdm(arg):
         tqdm_bars[0] = tqdm(range(n))
         tqdm_bars[0].set_description(message, refresh=False)
 
-    def _update_tqdm(arg, transform):
-        tqdm_bars[0].update(arg)
+    def _update_tqdm(arg):
+        tqdm_bars[0].update(float(arg))
 
     def _update_progress_bar(iter_num):
         "Updates tqdm from a JAX scan or loop"
         _ = jax.jax.lax.cond(
             iter_num == 0,
-            lambda _: host_callback.id_tap(_define_tqdm, None, result=iter_num),
-            lambda _: iter_num,
+            lambda _: callback(_define_tqdm, None),
+            lambda _: None,
             operand=None,
         )
 
         _ = jax.lax.cond(
             # update tqdm every multiple of `print_rate` except at the end
             (iter_num % print_rate == 0) & (iter_num != n - remainder),
-            lambda _: host_callback.id_tap(_update_tqdm, print_rate, result=iter_num),
-            lambda _: iter_num,
+            lambda _: callback(_update_tqdm, print_rate),
+            lambda _: None,
             operand=None,
         )
 
         _ = jax.lax.cond(
             # update tqdm by `remainder`
             iter_num == n - remainder,
-            lambda _: host_callback.id_tap(_update_tqdm, remainder, result=iter_num),
-            lambda _: iter_num,
+            lambda _: callback(_update_tqdm, remainder),
+            lambda _: None,
             operand=None,
         )
 
-    def _close_tqdm(arg, transform):
+    def _close_tqdm(arg):
         tqdm_bars[0].close()
+        return arg
 
     def close_tqdm(result, iter_num):
         return jax.lax.cond(
             iter_num == n - 1,
-            lambda _: host_callback.id_tap(_close_tqdm, None, result=result),
+            lambda _: io_callback(_close_tqdm, result),
             lambda _: result,
             operand=None,
         )
@@ -151,41 +153,41 @@ def new_while(cond_fun, body_fun, init_val):
         info = {'progval': 0, 'pbar': None}
         default_desc = f"Running {body_fun.__name__}"
 
-        def construct_tqdm(arg, transform):
+        def construct_tqdm(arg):
             _bar_format = "{l_bar}{bar}| {n:.3g}/{total_fmt} [{elapsed}<{remaining}, "
             _bar_format += "{rate_fmt}{postfix}]"
             info['pbar'] = tqdm(total=100, desc=f'{desc or default_desc}', bar_format=_bar_format)
 
-        def update_tqdm(arg, transform):
+        def update_tqdm(arg):
             error = errorfn(arg)
             errstart = info.setdefault('errstart', error)
             progress = max(100 * np.log(error / errstart) / np.log(tol / errstart) - info['progval'], 0)
             progress = min(100 - info['progval'], progress)
             if progress > 0:
                 info['progval'] += progress
-                info['pbar'].update(progress)
+                info['pbar'].update(float(progress))
 
-        def close_tqdm(arg, transform):
-            update_tqdm(arg, transform)
+        def close_tqdm(arg):
+            update_tqdm(arg)
             info['pbar'].close()
+            return False
 
         def newbody(ival):
             i, val = ival
             jax.lax.cond(
                 i % every == 0,
-                lambda _: host_callback.id_tap(update_tqdm, val, result=i),
-                lambda _: i,
+                lambda _: callback(update_tqdm, val),
+                lambda _: None,
                 operand=None,
             )
             return (i + 1, body_fun(val))
 
         def newcond(ival):
             i, val = ival
-            out = jax.lax.cond(cond_fun(val), lambda _: True,
-                               lambda _: host_callback.id_tap(close_tqdm, val, result=False), operand=None)
+            out = jax.lax.cond(cond_fun(val), lambda _: True, lambda _: io_callback(close_tqdm, val), operand=None)
             return out
 
-        host_callback.id_tap(construct_tqdm, None)
+        callback(construct_tqdm)
         _, val = jax.lax.while_loop(newcond, newbody, (0, init_val))
         return val
 
@@ -224,23 +226,23 @@ def construct_info(*_):
                 bar_format += "{rate_fmt}{postfix}]"
                 info['pbar'] = tqdm(total=100, desc=f'{desc or default_desc}', bar_format=bar_format)
 
-        def update_info(ival, _):
+        def update_info(ival):
             i, arg = ival
-            error = errorfn(arg)
+            error = float(errorfn(arg))
             info['errors'].append(error)
             if pbar:
                 errstart = info.setdefault('errstart', error)
                 howclose = np.log(error / errstart) / np.log(tol / errstart)
                 if max_iters is not None:
                     howclose = max((i + 1) / max_iters, howclose)
-                progress = min(100 - info['progval'], max(100 * howclose - info['progval'], 0))
+                progress = float(min(100 - info['progval'], max(100 * howclose - info['progval'], 0)))
                 if progress > 0:
                     info['progval'] += progress
                     info['pbar'].update(progress)
 
-        def close_info(arg, transform):
+        def close_info(arg):
             i, val = arg
-            update_info(arg, transform)
+            update_info(arg)
             info['iteration_time'] = (time.time() - info['iteration_time']) / (i + 1)
             if pbar:
                 info['pbar'].close()
@@ -254,8 +256,8 @@ def newbody(ival):
             i, val = ival
             jax.lax.cond(
                 i % every == 0,
-                lambda _: host_callback.id_tap(update_info, ival, result=i),
-                lambda _: i,
+                lambda _: callback(update_info, ival),
+                lambda _: None,
                 operand=None,
             )
             return (i + 1, body_fun(val))
@@ -269,9 +271,9 @@ def newcond(ival):
                 operand=None)
             return out
 
-        host_callback.id_tap(construct_info, None)
+        callback(construct_info, None)
         k, val = jax.lax.while_loop(newcond, newbody, (0, init_val))
-        host_callback.id_tap(close_info, (k, val))
+        callback(close_info, (k, val))
         return val
 
     return new_while, info

tests/test_linalg.py

@@ -32,7 +32,7 @@ def test_power_iteration(backend):
     A = xnp.diag(xnp.array([10., 9.75, 3., 0.1], dtype=dtype, device=None))
     B = lazify(A)
     soln = xnp.array(10., dtype=dtype, device=None)
-    tol, max_iter = 1e-5, 500
+    tol, max_iter = 1e-6, 500
     _, approx, _ = power_iteration(B, tol=tol, max_iter=max_iter, momentum=0.)
     rel_error = relative_error(soln, approx)
     assert rel_error < tol * 100