Repository: jax-ml/jax
Pull request: Add support for batched LU decomposition.
Fixes #655.

Problem statement:
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!

Hints from issue discussion:
I can add a simple implementation that batches by performing LU decompositions in a loop easily enough.

Commit message: Add test for QR batching.

Patch to review:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
index 75bde6362b37..f6b6838d6bab 100644
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -47,8 +47,7 @@ def eigh(x, lower=True, symmetrize_input=True):
 def lu(x): return lu_p.bind(x)
 
 def qr(x, full_matrices=True):
-  q, r = qr_p.bind(x, full_matrices=full_matrices)
-  return q, r
+  return qr_p.bind(x, full_matrices=full_matrices)
 
 def svd(x, full_matrices=True, compute_uv=True):
   s, u, v = svd_p.bind(x, full_matrices=full_matrices, compute_uv=compute_uv)
@@ -338,29 +337,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(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 +422,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 core.pack((q, r)), 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 +476,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(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
diff --git a/jaxlib/lapack.pyx b/jaxlib/lapack.pyx
index ac4ac8be0f13..3624aa65ad44 100644
--- a/jaxlib/lapack.pyx
+++ b/jaxlib/lapack.pyx
@@ -218,69 +218,89 @@ def jax_trsm(c, alpha, a, b, left_side=False, lower=False, trans_a=False,
 # ?getrf: LU decomposition
 
 cdef void lapack_sgetrf(void* out_tuple, void** data) nogil:
-  cdef int m = (<int32_t*>(data[0]))[0]
-  cdef int n = (<int32_t*>(data[1]))[0]
-  cdef const float* a_in = <float*>(data[2])
+  cdef int b = (<int32_t*>(data[0]))[0]
+  cdef int m = (<int32_t*>(data[1]))[0]
+  cdef int n = (<int32_t*>(data[2]))[0]
+  cdef const float* a_in = <float*>(data[3])
 
   cdef void** out = <void**>(out_tuple)
   cdef float* a_out = <float*>(out[0])
   cdef int* ipiv = <int*>(out[1])
   cdef int* info = <int*>(out[2])
   if a_out != a_in:
-    memcpy(a_out, a_in, m * n * sizeof(float))
+    memcpy(a_out, a_in, b * m * n * sizeof(float))
 
-  sgetrf(&m, &n, a_out, &m, ipiv, info)
+  for i in range(b):
+    sgetrf(&m, &n, a_out, &m, ipiv, info)
+    a_out += m * n
+    ipiv += min(m, n)
+    info += 1
 
 register_cpu_custom_call_target(b"lapack_sgetrf", <void*>(lapack_sgetrf))
 
 
 cdef void lapack_dgetrf(void* out_tuple, void** data) nogil:
-  cdef int m = (<int32_t*>(data[0]))[0]
-  cdef int n = (<int32_t*>(data[1]))[0]
-  cdef const double* a_in = <double*>(data[2])
+  cdef int b = (<int32_t*>(data[0]))[0]
+  cdef int m = (<int32_t*>(data[1]))[0]
+  cdef int n = (<int32_t*>(data[2]))[0]
+  cdef const double* a_in = <double*>(data[3])
 
   cdef void** out = <void**>(out_tuple)
   cdef double* a_out = <double*>(out[0])
   cdef int* ipiv = <int*>(out[1])
   cdef int* info = <int*>(out[2])
   if a_out != a_in:
-    memcpy(a_out, a_in, m * n * sizeof(double))
+    memcpy(a_out, a_in, b * m * n * sizeof(double))
 
-  dgetrf(&m, &n, a_out, &m, ipiv, info)
+  for i in range(b):
+    dgetrf(&m, &n, a_out, &m, ipiv, info)
+    a_out += m * n
+    ipiv += min(m, n)
+    info += 1
 
 register_cpu_custom_call_target(b"lapack_dgetrf", <void*>(lapack_dgetrf))
 
 
 cdef void lapack_cgetrf(void* out_tuple, void** data) nogil:
-  cdef int m = (<int32_t*>(data[0]))[0]
-  cdef int n = (<int32_t*>(data[1]))[0]
-  cdef const float complex* a_in = <float complex*>(data[2])
+  cdef int b = (<int32_t*>(data[0]))[0]
+  cdef int m = (<int32_t*>(data[1]))[0]
+  cdef int n = (<int32_t*>(data[2]))[0]
+  cdef const float complex* a_in = <float complex*>(data[3])
 
   cdef void** out = <void**>(out_tuple)
   cdef float complex* a_out = <float complex*>(out[0])
   cdef int* ipiv = <int*>(out[1])
   cdef int* info = <int*>(out[2])
   if a_out != a_in:
-    memcpy(a_out, a_in, m * n * sizeof(float complex))
+    memcpy(a_out, a_in, b * m * n * sizeof(float complex))
 
-  cgetrf(&m, &n, a_out, &m, ipiv, info)
+  for i in range(b):
+    cgetrf(&m, &n, a_out, &m, ipiv, info)
+    a_out += m * n
+    ipiv += min(m, n)
+    info += 1
 
 register_cpu_custom_call_target(b"lapack_cgetrf", <void*>(lapack_cgetrf))
 
 
 cdef void lapack_zgetrf(void* out_tuple, void** data) nogil:
-  cdef int m = (<int32_t*>(data[0]))[0]
-  cdef int n = (<int32_t*>(data[1]))[0]
-  cdef const double complex* a_in = <double complex*>(data[2])
+  cdef int b = (<int32_t*>(data[0]))[0]
+  cdef int m = (<int32_t*>(data[1]))[0]
+  cdef int n = (<int32_t*>(data[2]))[0]
+  cdef const double complex* a_in = <double complex*>(data[3])
 
   cdef void** out = <void**>(out_tuple)
   cdef double complex* a_out = <double complex*>(out[0])
   cdef int* ipiv = <int*>(out[1])
   cdef int* info = <int*>(out[2])
   if a_out != a_in:
-    memcpy(a_out, a_in, m * n * sizeof(double complex))
+    memcpy(a_out, a_in, b * m * n * sizeof(double complex))
 
-  zgetrf(&m, &n, a_out, &m, ipiv, info)
+  for i in range(b):
+    zgetrf(&m, &n, a_out, &m, ipiv, info)
+    a_out += m * n
+    ipiv += min(m, n)
+    info += 1
 
 register_cpu_custom_call_target(b"lapack_zgetrf", <void*>(lapack_zgetrf))
 
@@ -289,7 +309,15 @@ def jax_getrf(c, a):
 
   a_shape = c.GetShape(a)
   dtype = a_shape.element_type()
-  m, n = a_shape.dimensions()
+  dims = a_shape.dimensions()
+  assert len(dims) >= 2
+  m, n = dims[-2:]
+  batch_dims = tuple(dims[:-2])
+  num_bd = len(batch_dims)
+  b = 1
+  for d in batch_dims:
+    b *= d
+
   if dtype == np.float32:
     fn = b"lapack_sgetrf"
   elif dtype == np.float64:
@@ -303,16 +331,31 @@ def jax_getrf(c, a):
 
   return c.CustomCall(
       fn,
-      operands=(c.ConstantS32Scalar(m), c.ConstantS32Scalar(n), a),
+      operands=(
+        c.ConstantS32Scalar(b),
+        c.ConstantS32Scalar(m),
+        c.ConstantS32Scalar(n),
+        a),
       shape_with_layout=Shape.tuple_shape((
-          Shape.array_shape(dtype, (m, n), (0, 1)),
-          Shape.array_shape(np.int32, (min(m, n),), (0,)),
-          Shape.array_shape(np.int32, (), ()),
+          Shape.array_shape(
+            dtype,
+            batch_dims + (m, n),
+            (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))),
+          Shape.array_shape(
+            np.int32,
+            batch_dims + (min(m, n),),
+            tuple(range(num_bd, -1, -1))),
+          Shape.array_shape(np.int32, batch_dims,
+            tuple(range(num_bd - 1, -1, -1))),
       )),
       operand_shapes_with_layout=(
           Shape.array_shape(np.int32, (), ()),
           Shape.array_shape(np.int32, (), ()),
-          Shape.array_shape(dtype, (m, n), (0, 1)),
+          Shape.array_shape(np.int32, (), ()),
+          Shape.array_shape(
+            dtype,
+            batch_dims + (m, n),
+            (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))),
       ))
 
 
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index df8a89161d0a..d00a013c509c 100644
--- 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,21 @@ 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):
+    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(


Code review comments:
diff_hunk: @@ -47,8 +47,7 @@ def eigh(x, lower=True, symmetrize_input=True):
 def lu(x): return lu_p.bind(x)
 
 def qr(x, full_matrices=True):
-  q, r = qr_p.bind(x, full_matrices=full_matrices)
-  return q, r
+  return qr_p.bind(x, full_matrices=full_matrices)
path: jax/lax_linalg.py
text: @user1:
After this change, we have a traceable (user-callable wrapper around a primitive's bind rule) returning a JaxTuple. We generally don't want that; I think we should revert this change.

@author:
Done. And made the opposite change everywhere else in this file that returns a tuple (e.g. `eigh`)

diff_hunk: @@ -427,7 +426,8 @@ def qr_batching_rule(batched_args, batch_dims, full_matrices):
   x, = batched_args
   bd, = batch_dims
   x = batching.bdim_at_front(x, bd)
-  return qr(x, full_matrices=full_matrices), 0
+  q, r = qr(x, full_matrices=full_matrices)
+  return core.pack((q, r)), 0
path: jax/lax_linalg.py
text: @user1:
Another way to write this would be

```python
return qr_p.bind(x, full_matrices=full_matrices), 0
```

@author:
Done.