Repository: jax-ml/jax
Pull request: Eliminate redundancies in jax.scipy.linalg.expm
## Summary

This PR introduces a minor change to the `expm` function:
- Currently, given the matrix `A`, `expm` computes all Pade approximant numerators/denominators, then selects which one to use based on the norm of the matrix.
- This PR modifies this to only compute the relevant Pade approximant numerator/denominator by using `switch` conditioned on the norm of the matrix `A`.

Closes #5647

## Speed testing

I unfortunately don't currently have access to a GPU, but have run some CPU speed tests comparing the code in this PR to the current release. Testing in 64 bit mode with complex matrices of dimensions `[5, 10, 50, 100, 500, 1000, 5000]`, the ratios by dimension of "release run time" / "PR run time" for `expm` (once compiled) are:
```
[1.0794702 , 3.63815789, 3.41229117, 3.38029568, 1.68847097, 1.54355993, 1.61012301]
```
These are not the most rigorous (I'm just generating random matrices and measuring time by sandwiching `time()` calls around a single exponentiation for each dimension), but seem pretty reasonable based on the number of matrix multiplications this is potentially saving (which varies based on matrix norm).

I've attached the jupyter notebook I've run the tests in [jax_expm_speed_test.ipynb.zip](https://github.com/google/jax/files/7898241/jax_expm_speed_test.ipynb.zip), which include a cell of the saved list of execution times for JAX version (release and this PR). 

## Going forward

I'm trying to finally get GPU access in my org and will reply with GPU speed comparisons here when possible. In the meantime, please let me know if there's anything you'd like to see in terms of speed comparisons.

Problem statement:
Redundant computation in expm?
Hi all:

It looks to me there are redundant calculations in the `expm` function, but wanted to ask to see if I'm missing something (perhaps something clever is happening in compilation). 

E.g., for `float64` and `complex128`, the numerator/denominator for the Pade approximant is generated as follows:

```
if A.dtype == 'float64' or A.dtype == 'complex128':
   U3, V3 = _pade3(A)
   U5, V5 = _pade5(A)
   U7, V7 = _pade7(A)
   U9, V9 = _pade9(A)
   maxnorm = 5.371920351148152
   n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
   A = A / 2**n_squarings
   U13, V13 = _pade13(A)
   conds=jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
                    9.504178996162932e-001, 2.097847961257068e+000])
   U = jnp.select((A_L1<conds), (U3, U5, U7, U9), U13)
   V = jnp.select((A_L1<conds), (V3, V5, V7, V9), V13)
```

The algorithm being used (I believe by Al-Mohy and Higham, same as in `scipy`) prescribes choosing between 1 of 5 different Pade approximations (of various orders) based on the norm of `A`. In the above, it appears all potential Pade approximations are being computed, and only afterwards is one chosen based on the norm of `A`. Hence, the computations for the others are discarded and ultimately unnecessary.

It seems this function could benefit from changing the above logic to use nested `lax.cond` statements (or perhaps `switch`?) to avoid computing all possible versions. To motivate this: I quickly tried modifying the `expm` function to simply always do the `U13, V13` version of of the computation (the most expensive) and not compute any of the lower order versions, and observed roughly a factor of 2 speed increase in compiled execution. My speed tests weren't rigorous - I didn't test it for a variety of norms - but I believe this should speed things up in all cases (again, unless I'm missing something with lax behaviour). Edit: This is on CPU.

Thoughts?

P.s. In some of my use cases of jax, the most expensive piece is a ton of matrix exponentials, so I'm very interested in potential speed gains for this function :D.

Hints from issue discussion:
This sounds right to me! Indeed using `jnp.select` is wasting compute, and `cond`/`switch` could save it. I don't remember any context around why `expm` was written this way (@zhangqiaorjc do you remember?) but it may just have been for simplicity. In that case, what you describe would be a welcome upgrade!
I guess that on GPU it _might_ be better to compute everything rather than pass control flow back to the host. But that's really an optimization problem for XLA GPU, not something we should worry about in JAX.

In general, matrix-matrix multiplication (L3 BLAS) seems like the level of computation where we should be using `cond` rather than `where`.
Does `cond` always pass control flow back to the host on GPU? That sounds plausible, though also it seems like at least in some cases XLA:GPU could avoid that. (At the very least, it could decide to optimize a `cond` into a `select` if needed!) On TPU (where control flow can stay on the device) and CPU (where there's no host/device separation) it seems profitable.

Good rule of thumb!
If there are no objections, given how small this is, and my personal interest in seeing it go through asap, I'm willing to implement this change. My plan would be to just change it to use nested `cond` statements. In a PR I can show some local CPU benchmarks.
@DanPuzzuoli Feel free to send in a PR. I don't think anyone from JAX team is currently working on this. Thanks!

Commit message: changed to use argwhere

Patch to review:
diff --git a/jax/_src/scipy/linalg.py b/jax/_src/scipy/linalg.py
index bd710892a3e8..7f0212b74028 100644
--- a/jax/_src/scipy/linalg.py
+++ b/jax/_src/scipy/linalg.py
@@ -289,28 +289,20 @@ def _calc_P_Q(A):
   A_L1 = np_linalg.norm(A,1)
   n_squarings = 0
   if A.dtype == 'float64' or A.dtype == 'complex128':
-   U3, V3 = _pade3(A)
-   U5, V5 = _pade5(A)
-   U7, V7 = _pade7(A)
-   U9, V9 = _pade9(A)
    maxnorm = 5.371920351148152
    n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
    A = A / 2**n_squarings
-   U13, V13 = _pade13(A)
    conds = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
-                      9.504178996162932e-001, 2.097847961257068e+000])
-   U = jnp.select((A_L1 < conds), (U3, U5, U7, U9), U13)
-   V = jnp.select((A_L1 < conds), (V3, V5, V7, V9), V13)
+                      9.504178996162932e-001, 2.097847961257068e+000, jnp.inf])
+   idx = jnp.argwhere(A_L1 < conds, size=1).squeeze()
+   U, V = lax.switch(idx, [_pade3, _pade5, _pade7, _pade9, _pade13], A)
   elif A.dtype == 'float32' or A.dtype == 'complex64':
-    U3,V3 = _pade3(A)
-    U5,V5 = _pade5(A)
     maxnorm = 3.925724783138660
     n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
     A = A / 2**n_squarings
-    U7,V7 = _pade7(A)
-    conds = jnp.array([4.258730016922831e-001, 1.880152677804762e+000])
-    U = jnp.select((A_L1 < conds), (U3, U5), U7)
-    V = jnp.select((A_L1 < conds), (V3, V5), V7)
+    conds = jnp.array([4.258730016922831e-001, 1.880152677804762e+000, jnp.inf])
+    idx = jnp.argwhere(A_L1 < conds, size=1).squeeze()
+    U, V = lax.switch(idx, [_pade3, _pade5, _pade7], A)
   else:
     raise TypeError("A.dtype={} is not supported.".format(A.dtype))
   P = U + V  # p_m(A) : numerator


Code review comments:
diff_hunk: @@ -289,28 +289,20 @@ def _calc_P_Q(A):
   A_L1 = np_linalg.norm(A,1)
   n_squarings = 0
   if A.dtype == 'float64' or A.dtype == 'complex128':
-   U3, V3 = _pade3(A)
-   U5, V5 = _pade5(A)
-   U7, V7 = _pade7(A)
-   U9, V9 = _pade9(A)
    maxnorm = 5.371920351148152
    n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
    A = A / 2**n_squarings
-   U13, V13 = _pade13(A)
    conds = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
-                      9.504178996162932e-001, 2.097847961257068e+000])
-   U = jnp.select((A_L1 < conds), (U3, U5, U7, U9), U13)
-   V = jnp.select((A_L1 < conds), (V3, V5, V7, V9), V13)
+                      9.504178996162932e-001, 2.097847961257068e+000, jnp.inf])
+   idx = jnp.argwhere(A_L1 < conds, size=1).squeeze()
path: jax/_src/scipy/linalg.py
text: @user1:
Ah, this is your use case. Then you might also want to try jnp.digitize.

@author:
Nice yes that is cleaner.

@user1:
But be aware that `jnp.digitize` is `searchsorted`-based, which could have different performance characteristics from the `nonzero`-based `jnp.argwhere`. I have absolutely no idea which one should/would be faster in this particular case, only through benchmark can one tell I guess.

@author:
Ah I see - thanks. My naive expectation is that the matrix operations would dominate in this function anyway, especially given how few options either `digitize` or `argwhere` need to choose from.

@user2:
I think this would give the same result, and be more efficient than either `argwhere` or `digitize` in their current implementations:
```
idx = jnp.argmax(A_L1 < conds)
```
Explanation: `argmax` returns the index of the first appearance of the maximum value in an array, and True is the maximum possible boolean value.

@author:
Okay so I've just run the following code:

```
conds = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
                      9.504178996162932e-001, 2.097847961257068e+000, jnp.inf])
conds2 = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
                      9.504178996162932e-001, 2.097847961257068e+000])
```
and in different cells, for various values of `x`, I'm seeing:
```
%timeit jnp.argmax(x < conds).block_until_ready()
```
CPU (local): 178 microseconds
GPU (colab): 610 microseconds
```
%timeit jnp.argwhere(x < conds, size=1).squeeze().block_until_ready()
```
CPU (local): 1.89 ms
GPU (colab): 5.09 ms
```
%timeit jnp.digitize(x, conds2).block_until_ready()
```
CPU (local): 5.8 microseconds
GPU (colab): 270 microseconds

The above values were for `x=3.` but the speed rankings seem to be the same regardless of `x`. Seems like for this case `digitize` is the clear winner.

Going to revert this branch back to the digitize version - let me know if you think there are any issues with this.

@user2:
OK, sounds good!