Repository: google/jax

Issue:
Add impl of `scipy.linalg.lu_solve`
This is used by the Glow generative model to parameterize an invertible 1x1 convolution.

A TF implementation in terms of triangular solve is here:
https://github.com/tensorflow/probability/blob/master/tensorflow_probability/python/math/linalg.py#L306


Discussion:
Pretty much exactly the same functionality is already available via `jax.numpy.linalg.solve`, but I guess there's no harm in adding the `scipy` variant too. (It supports an additional transpose or conjugate transpose and takes its input prefactored.)

Resolving patch:
diff --git a/jax/scipy/linalg.py b/jax/scipy/linalg.py
--- a/jax/scipy/linalg.py
+++ b/jax/scipy/linalg.py
@@ -16,10 +16,11 @@
 from __future__ import division
 from __future__ import print_function
 
-import warnings
+from functools import partial
 
 import scipy.linalg
 
+from jax import jit
 from .. import lax
 from .. import lax_linalg
 from ..numpy.lax_numpy import _wraps
@@ -117,6 +118,46 @@ def lu_factor(a, overwrite_a=False, check_finite=True):
   a = np_linalg._promote_arg_dtypes(np.asarray(a))
   return lax_linalg.lu(a)
 
+@partial(jit, static_argnums=(3,))
+def _lu_solve(lu, pivots, b, trans):
+  lu_shape = np.shape(lu)
+  b_shape = np.shape(b)
+  if len(lu_shape) != 2 or lu_shape[0] != lu_shape[1]:
+    raise ValueError("LU decomposition must be a square matrix, got shape {}"
+                     .format(lu_shape))
+  if len(b_shape) < 1:
+    raise ValueError("b matrix must have rank >= 1, got shape {}"
+                     .format(b_shape))
+
+  if b_shape[0] != lu_shape[0]:
+    raise ValueError("Dimension of LU decomposition matrix (shape {}) must "
+                     "match leading axis of b array (shape {})"
+                     .format(lu_shape, b_shape))
+  m = lu_shape[0]
+  permutation = lax_linalg.lu_pivots_to_permutation(np.array(pivots), m)
+  x = np.reshape(b, (m, -1))
+  if trans == 0:
+    x = x[permutation, :]
+    x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=True,
+                                    unit_diagonal=True)
+    x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=False)
+  elif trans == 1 or trans == 2:
+    conj = trans == 2
+    x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=False,
+                                    transpose_a=True, conjugate_a=conj)
+    x = lax_linalg.triangular_solve(lu, x, left_side=True, lower=True,
+                                    unit_diagonal=True, transpose_a=True,
+                                    conjugate_a=conj)
+    x = x[np.argsort(permutation), :]
+  else:
+    raise ValueError("'trans' value must be 0, 1, or 2, got {}".format(trans))
+  return lax.reshape(x, b_shape)
+
+@_wraps(scipy.linalg.lu_solve)
+def lu_solve(lu_and_piv, b, trans=0, overwrite_b=False, check_finite=True):
+  del overwrite_b, check_finite
+  lu, pivots = lu_and_piv
+  return _lu_solve(lu, pivots, b, trans)
 
 @_wraps(scipy.linalg.lu)
 def lu(a, permute_l=False, overwrite_a=False, check_finite=True):


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -585,6 +585,36 @@ def testLuFactor(self, n, dtype, rng):
     self.assertAllClose(x, onp.matmul(l, u), check_dtypes=True, rtol=1e-3)
     self._CompileAndCheck(jsp.linalg.lu_factor, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name":
+       "_lhs={}_rhs={}_trans={}".format(
+           jtu.format_shape_dtype_string(lhs_shape, dtype),
+           jtu.format_shape_dtype_string(rhs_shape, dtype),
+           trans),
+       "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
+       "trans": trans, "rng": rng}
+      for lhs_shape, rhs_shape in [
+          ((1, 1), (1, 1)),
+          ((4, 4), (4,)),
+          ((8, 8), (8, 4, 2)),
+      ]
+      for trans in [0, 1, 2]
+      for dtype in float_types + complex_types
+      for rng in [jtu.rand_default()]))
+  def testLuSolve(self, lhs_shape, rhs_shape, dtype, trans, rng):
+    _skip_if_unsupported_type(dtype)
+    osp_fun = lambda lu, piv, rhs: osp.linalg.lu_solve((lu, piv), rhs, trans=trans)
+    jsp_fun = lambda lu, piv, rhs: jsp.linalg.lu_solve((lu, piv), rhs, trans=trans)
+
+    def args_maker():
+      a = rng(lhs_shape, dtype)
+      lu, piv = osp.linalg.lu_factor(a)
+      return [lu, piv, rng(rhs_shape, dtype)]
+
+    self._CheckAgainstNumpy(osp_fun, jsp_fun, args_maker,
+                            check_dtypes=True, tol=1e-3)
+    self._CompileAndCheck(jsp_fun, args_maker, check_dtypes=True)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name":
        "_lhs={}_rhs={}_sym_pos={}_lower={}".format(