Repository: google/jax

Issue:
Implement vmap/batching support for LU decomposition 
I'm trying to solve linear equations Ax = b with vmap:

```python
import jax.numpy as np
import jax.random as random
from jax import vmap

key = random.PRNGKey(0)

A = random.normal(key, shape=(100, 2, 2))
x = random.normal(key, shape=(100, 2))
b = vmap(lambda A, x: A @ x)(A, x) # works! shape (100, 2)
x_hat = vmap(lambda A, x: np.linalg.solve(A, x))(A, x)
```

Here, using `np.linalg.solve` results in the following error message:

```
NotImplementedError: Batching rule for 'lu' not implemented
```

Opening this issue to request batching support for LU decomposition. Thanks in advance!


Discussion:
I can add a simple implementation that batches by performing LU decompositions in a loop easily enough.

Resolving patch:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -42,9 +42,12 @@ def cholesky(x, symmetrize_input=True):
 def eigh(x, lower=True, symmetrize_input=True):
   if symmetrize_input:
     x = symmetrize(x)
-  return eigh_p.bind(x, lower=lower)
+  v, w = eigh_p.bind(x, lower=lower)
+  return v, w
 
-def lu(x): return lu_p.bind(x)
+def lu(x):
+  lu, pivots = lu_p.bind(x)
+  return lu, pivots
 
 def qr(x, full_matrices=True):
   q, r = qr_p.bind(x, full_matrices=full_matrices)
@@ -338,29 +341,31 @@ def lu_jvp_rule(primals, tangents):
   return core.pack((lu, pivots)), ad.TangentTuple((lu_dot, ad_util.zero))
 
 
+def lu_batching_rule(batched_args, batch_dims):
+  x, = batched_args
+  bd, = batch_dims
+  x = batching.bdim_at_front(x, bd)
+  return lu_p.bind(x), 0
+
+
 lu_p = Primitive('lu')
 lu_p.def_impl(lu_impl)
 lu_p.def_abstract_eval(lu_abstract_eval)
 xla.translations[lu_p] = lu_translation_rule
 ad.primitive_jvps[lu_p] = lu_jvp_rule
+batching.primitive_batchers[lu_p] = lu_batching_rule
 
 def lu_cpu_translation_rule(c, operand):
   shape = c.GetShape(operand)
   dtype = shape.element_type().type
-  if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
-    out = lapack.jax_getrf(c, operand)
-    lu = c.GetTupleElement(out, 0)
-    # Subtract 1 from the pivot to get 0-based indices.
-    pivot = c.Sub(c.GetTupleElement(out, 1), c.ConstantS32Scalar(1))
-    # Throw away the `info` value, because we have no way to report errors.
-    return c.Tuple(lu, pivot)
-  else:
-    raise NotImplementedError("Only unbatched LU decomposition is implemented")
+  out = lapack.jax_getrf(c, operand)
+  lu = c.GetTupleElement(out, 0)
+  # Subtract 1 from the pivot to get 0-based indices.
+  pivot = c.Sub(c.GetTupleElement(out, 1), c.ConstantS32Scalar(1))
+  # Throw away the `info` value, because we have no way to report errors.
+  return c.Tuple(lu, pivot)
 
-# TODO(phawkins): The hasattr() test here is to avoid incompatibilities between
-# jax and an older jaxlib. Remove after a jaxlib release includes jax_getrf.
-if hasattr(lapack, "jax_getrf"):
-  xla.backend_specific_translations['cpu'][lu_p] = lu_cpu_translation_rule
+xla.backend_specific_translations['cpu'][lu_p] = lu_cpu_translation_rule
 
 
 def lu_pivots_to_permutation(swaps, k):
@@ -421,11 +426,19 @@ def qr_jvp_rule(primals, tangents, full_matrices):
   dr = np.matmul(qt_dx_rinv - domega, r)
   return core.pack((q, r)), core.pack((dq, dr))
 
+def qr_batching_rule(batched_args, batch_dims, full_matrices):
+  x, = batched_args
+  bd, = batch_dims
+  x = batching.bdim_at_front(x, bd)
+  q, r = qr(x, full_matrices=full_matrices)
+  return qr_p.bind(x, full_matrices=full_matrices), 0
+
 qr_p = Primitive('qr')
 qr_p.def_impl(qr_impl)
 qr_p.def_abstract_eval(qr_abstract_eval)
 xla.translations[qr_p] = qr_translation_rule
 ad.primitive_jvps[qr_p] = qr_jvp_rule
+batching.primitive_batchers[qr_p] = qr_batching_rule
 
 
 # Singular value decomposition
@@ -467,8 +480,15 @@ def svd_cpu_translation_rule(c, operand, full_matrices, compute_uv):
     raise NotImplementedError(
         "Only unbatched singular value decomposition is implemented on CPU")
 
+def svd_batching_rule(batched_args, batch_dims, full_matrices, compute_uv):
+  x, = batched_args
+  bd, = batch_dims
+  x = batching.bdim_at_front(x, bd)
+  return svd_p.bind(x, full_matrices=full_matrices, compute_uv=compute_uv), 0
+
 svd_p = Primitive('svd')
 svd_p.def_impl(svd_impl)
 svd_p.def_abstract_eval(svd_abstract_eval)
 xla.translations[svd_p] = svd_translation_rule
 xla.backend_specific_translations['cpu'][svd_p] = svd_cpu_translation_rule
+batching.primitive_batchers[svd_p] = svd_batching_rule


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -27,7 +27,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 
-from jax import jvp
+from jax import jvp, vmap
 from jax import numpy as np
 from jax import scipy as jsp
 from jax import test_util as jtu
@@ -340,6 +340,15 @@ def compare_orthogonal(q1, q2):
     if not full_matrices and m >= n:
         jtu.check_jvp(np.linalg.qr, partial(jvp, np.linalg.qr), (a,))
 
+  @jtu.skip_on_devices("gpu", "tpu")
+  def testQrBatching(self):
+    shape = (10, 4, 5)
+    dtype = np.float32
+    rng = jtu.rand_default()
+    args = rng(shape, np.float32)
+    qs, rs = vmap(jsp.linalg.qr)(args)
+    self.assertTrue(onp.all(onp.linalg.norm(args - onp.matmul(qs, rs)) < 1e-3))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name":
        "_lhs={}_rhs={}".format(
@@ -419,6 +428,22 @@ def testLuGrad(self, shape, dtype, rng):
 
     jtu.check_grads(jsp.linalg.lu, (a,), 2, rtol=1e-1)
 
+  @jtu.skip_on_devices("gpu", "tpu")
+  def testLuBatching(self):
+    self.skipTest("Test disabled until Jaxlib 0.1.14 is released")
+    shape = (4, 5)
+    dtype = np.float32
+    rng = jtu.rand_default()
+    args = [rng(shape, np.float32) for _ in range(10)]
+    expected = list(osp.linalg.lu(x) for x in args)
+    ps = onp.stack([out[0] for out in expected])
+    ls = onp.stack([out[1] for out in expected])
+    us = onp.stack([out[2] for out in expected])
+
+    actual_ps, actual_ls, actual_us = vmap(jsp.linalg.lu)(np.stack(args))
+    self.assertAllClose(ps, actual_ps, check_dtypes=True)
+    self.assertAllClose(ls, actual_ls, check_dtypes=True)
+    self.assertAllClose(us, actual_us, check_dtypes=True)
 
   # TODO(phawkins): enable when there is an LU implementation for GPU/TPU.
   @parameterized.named_parameters(jtu.cases_from_list(