Repository: google/jax

Issue:
Direct gradient rules to speed-up numpy.linalg.solve
It can be [significantly more efficient](https://www.semanticscholar.org/paper/Structured-higher-order-algorithmic-differentiation-Walter/e52d834dac4fdfdc4328105e29ae657c3fa51467), both in runtime and memory, to directly define derivative rules for higher level linear algebra operations like `solve` rather than the constituent operations (factorization and triangular solve).

For large matrices (e.g., 500x500 on CPUs), my microbenchmark shows that we can get a 3-4x speed-up for general purpose and symmetric solves:
```python
from functools import partial
import jax.scipy as jsp
from jax import lax
import jax.numpy as np
import numpy as onp
import jax

def positive_definite_solve(a, b):
  factors = jsp.linalg.cho_factor(a)
  def solve(matvec, x):
    return jsp.linalg.cho_solve(factors, x)
  matvec = partial(np.dot, a)
  return lax.custom_linear_solve(matvec, b, solve, symmetric=True)

def linear_solve(a, b):
  a_factors = jsp.linalg.lu_factor(a)
  def solve(matvec, x):
    return jsp.linalg.lu_solve(a_factors, x)
  def transpose_solve(vecmat, x):
    return jsp.linalg.lu_solve(a_factors, x, trans=1)
  matvec = partial(np.dot, a)
  return lax.custom_linear_solve(matvec, b, solve, transpose_solve)

def loss(solve):
  def f(a, b):
    return solve(a, b).sum()
  return f

rs = onp.random.RandomState(0)
a = rs.randn(500, 500)
a = jax.device_put(a.T @ a + 0.1 * np.eye(500))
b = jax.device_put(rs.randn(500))

# general purpose solve
# current
grad = jax.jit(jax.grad(loss(np.linalg.solve)))
%timeit jax.device_get(grad(a, b))
# 33.8 ms per loop
# new
grad = jax.jit(jax.grad(loss(linear_solve)))
%timeit jax.device_get(grad(a, b))
# 10.1 ms per loop

# positive definite solve
# current
grad = jax.jit(jax.grad(loss(partial(jsp.linalg.solve, sym_pos=True))))
%timeit jax.device_get(grad(a, b))
# 23.7 ms per loop
# new
grad = jax.jit(jax.grad(loss(positive_definite_solve)))
%timeit jax.device_get(grad(a, b))
# 4.8 ms per loop
```

Unfortunately, we can't just use these prototype implementations internally in JAX, for two reasons:
1. ~~`custom_linear_solve` (like custom transforms in general) doesn't work with batching yet (https://github.com/google/jax/issues/1249).~~ This was solved by #2099.
2. ~~We do an extra optimization in `triangular_solve_jvp_rule_a` for the case of solving many right-hand-sides at the same time with the same left-hand side (https://github.com/google/jax/pull/1466). This new gradient rule here doesn't handle this yet.~~ Update: in practice, I don't think this optimization actually matters -- it's the difference between `n*m*m+m*m*m` time vs `2*m*m*m` time.
3. We need to support multiple right-hand-side arguments. After https://github.com/google/jax/pull/2138, we'll be able to do this by batching `custom_linear_solve`.


Discussion:
Thanks to @mattjj  for pointing out that LU solve has the `trans` argument, which means we use a single factorization for both forward and reverse calculations to speed up solves on all types of matrices.

Resolving patch:
diff --git a/jax/numpy/linalg.py b/jax/numpy/linalg.py
--- a/jax/numpy/linalg.py
+++ b/jax/numpy/linalg.py
@@ -27,6 +27,7 @@
 from .. import dtypes
 from .lax_numpy import _not_implemented
 from .lax_numpy import _wraps
+from .vectorize import vectorize
 from . import lax_numpy as np
 from ..api import custom_transforms, defjvp
 from ..util import get_module_functions
@@ -328,16 +329,38 @@ def qr(a, mode="reduced"):
   return q, r
 
 
-@_wraps(onp.linalg.solve)
-@jit
-def solve(a, b):
-  a, b = _promote_arg_dtypes(np.asarray(a), np.asarray(b))
+def _check_solve_shapes(a, b):
   if not (a.ndim >= 2 and a.shape[-1] == a.shape[-2] and b.ndim >= 1):
     msg = ("The arguments to solve must have shapes a=[..., m, m] and "
            "b=[..., m, k] or b=[..., m]; got a={} and b={}")
     raise ValueError(msg.format(a.shape, b.shape))
-  lu, pivots = lax_linalg.lu(a)
-  return lax_linalg.lu_solve(lu, pivots, b, trans=0)
+
+
+@partial(vectorize, signature='(n,m),(m)->(n)')
+def _matvec_multiply(a, b):
+  return np.dot(a, b, precision=lax.Precision.HIGHEST)
+
+
+@_wraps(onp.linalg.solve)
+@jit
+def solve(a, b):
+  a, b = _promote_arg_dtypes(np.asarray(a), np.asarray(b))
+  _check_solve_shapes(a, b)
+
+  # With custom_linear_solve, we can reuse the same factorization when
+  # computing sensitivities. This is considerably faster.
+  lu, pivots = lax.stop_gradient(lax_linalg.lu)(a)
+  custom_solve = partial(
+      lax.custom_linear_solve,
+      lambda x: _matvec_multiply(a, x),
+      solve=lambda _, x: lax_linalg.lu_solve(lu, pivots, x, trans=0),
+      transpose_solve=lambda _, x: lax_linalg.lu_solve(lu, pivots, x, trans=1))
+  if a.ndim == b.ndim + 1:
+    # b.shape == [..., m]
+    return custom_solve(b)
+  else:
+    # b.shape == [..., m, k]
+    return vmap(custom_solve, b.ndim - 1, max(a.ndim, b.ndim) - 1)(b)
 
 
 for func in get_module_functions(onp.linalg):
diff --git a/jax/scipy/linalg.py b/jax/scipy/linalg.py
--- a/jax/scipy/linalg.py
+++ b/jax/scipy/linalg.py
@@ -18,7 +18,7 @@
 import scipy.linalg
 import textwrap
 
-from jax import jit
+from jax import jit, vmap
 from .. import lax
 from .. import lax_linalg
 from ..numpy.lax_numpy import _wraps
@@ -45,24 +45,16 @@ def cho_factor(a, lower=False, overwrite_a=False, check_finite=True):
 @partial(jit, static_argnums=(2,))
 def _cho_solve(c, b, lower):
   c, b = np_linalg._promote_arg_dtypes(np.asarray(c), np.asarray(b))
-  c_shape = np.shape(c)
-  b_shape = np.shape(b)
-  c_ndims = len(c_shape)
-  b_ndims = len(b_shape)
-  if not (c_ndims >= 2 and c_shape[-1] == c_shape[-2] and
-          (c_ndims == b_ndims or c_ndims == b_ndims + 1)):
-    msg = ("The arguments to solve must have shapes a=[..., m, m] and "
-           "b=[..., m, k] or b=[..., m]; got a={} and b={}")
-    raise ValueError(msg.format(c_shape, b_shape))
-
+  np_linalg._check_solve_shapes(c, b)
   # TODO(phawkins): triangular_solve only supports matrices on the RHS, so we
   # add a dummy dimension. Extend it to support vectors and simplify this.
-  b = b if c_ndims == b_ndims else b[..., None]
+  rhs_vector = c.ndim == b.ndim + 1
+  b = b[..., np.newaxis] if rhs_vector else b
   b = lax_linalg.triangular_solve(c, b, left_side=True, lower=lower,
                                   transpose_a=not lower, conjugate_a=not lower)
   b = lax_linalg.triangular_solve(c, b, left_side=True, lower=lower,
                                   transpose_a=lower, conjugate_a=lower)
-  return b[..., 0] if c_ndims != b_ndims else b
+  return b[..., 0] if rhs_vector else b
 
 @_wraps(scipy.linalg.cho_solve, update_doc=False)
 def cho_solve(c_and_lower, b, overwrite_b=False, check_finite=True):
@@ -170,13 +162,30 @@ def qr(a, overwrite_a=False, lwork=None, mode="full", pivoting=False,
   del overwrite_a, lwork, check_finite
   return _qr(a, mode, pivoting)
 
+
 @partial(jit, static_argnums=(2, 3))
 def _solve(a, b, sym_pos, lower):
   if not sym_pos:
     return np_linalg.solve(a, b)
 
   a, b = np_linalg._promote_arg_dtypes(np.asarray(a), np.asarray(b))
-  return cho_solve(cho_factor(a, lower=lower), b)
+  np_linalg._check_solve_shapes(a, b)
+
+  # With custom_linear_solve, we can reuse the same factorization when
+  # computing sensitivities. This is considerably faster.
+  factors = lax.stop_gradient(cho_factor)(a, lower=lower)
+  custom_solve = partial(
+      lax.custom_linear_solve,
+      lambda x: np_linalg._matvec_multiply(a, x),
+      solve=lambda _, x: cho_solve(factors, x),
+      symmetric=True)
+  if a.ndim == b.ndim + 1:
+    # b.shape == [..., m]
+    return custom_solve(b)
+  else:
+    # b.shape == [..., m, k]
+    return vmap(custom_solve, b.ndim - 1, max(a.ndim, b.ndim) - 1)(b)
+
 
 @_wraps(scipy.linalg.solve)
 def solve(a, b, sym_pos=False, lower=False, overwrite_a=False, overwrite_b=False,


Test patch: