Repository: google/jax

Issue:
add batching (vmap) rule for eigh



Discussion:

Resolving patch:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -143,14 +143,8 @@ def eigh_abstract_eval(operand, lower):
   return core.AbstractTuple((v, w))
 
 def eigh_cpu_translation_rule(c, operand, lower):
-  shape = c.GetShape(operand)
-  dtype = shape.element_type().type
-  if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
-    out = lapack.jax_syevd(c, operand, lower=lower)
-    return c.Tuple(c.GetTupleElement(out, 0), c.GetTupleElement(out, 1))
-  else:
-    raise NotImplementedError(
-        "Only unbatched eigendecomposition is implemented on CPU")
+  out = lapack.jax_syevd(c, operand, lower=lower)
+  return c.Tuple(c.GetTupleElement(out, 0), c.GetTupleElement(out, 1))
 
 def eigh_jvp_rule(primals, tangents, lower):
   # Derivative for eigh in the simplest case of distinct eigenvalues.
@@ -176,12 +170,19 @@ def eigh_jvp_rule(primals, tangents, lower):
   dw = np.diagonal(vdag_adot_v)
   return core.pack((v, w)), core.pack((dv, dw))
 
+def eigh_batching_rule(batched_args, batch_dims, lower):
+  x, = batched_args
+  bd, = batch_dims
+  x = batching.bdim_at_front(x, bd)
+  return eigh_p.bind(x, lower=lower), 0
+
 eigh_p = Primitive('eigh')
 eigh_p.def_impl(eigh_impl)
 eigh_p.def_abstract_eval(eigh_abstract_eval)
 xla.translations[eigh_p] = eigh_translation_rule
 ad.primitive_jvps[eigh_p] = eigh_jvp_rule
 xla.backend_specific_translations['cpu'][eigh_p] = eigh_cpu_translation_rule
+batching.primitive_batchers[eigh_p] = eigh_batching_rule
 
 
 
@@ -430,7 +431,6 @@ 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')


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -204,6 +204,23 @@ def testEighGradVectorComplex(self, shape, dtype, rng, lower, eps):
       onp.linalg.norm(onp.abs(new_w*(v+dv)), axis=0)
     ) < RTOL
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name":
+       "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
+       "shape": shape, "dtype": dtype, "rng": rng}
+      for shape in [(1, 1), (4, 4), (5, 5)]
+      for dtype in float_types() + complex_types()
+      for rng in [jtu.rand_default()]))
+  @jtu.skip_on_devices("gpu", "tpu")
+  def testEighBatching(self, shape, dtype, rng):
+    self.skipTest("Test disabled until Jaxlib 0.1.15 is released") # TODO(phawkins)
+    shape = (10,) + shape
+    args = rng(shape, dtype)
+    args = (args + onp.conj(T(args))) / 2
+    ws, vs = vmap(jsp.linalg.eigh)(args)
+    self.assertTrue(onp.all(onp.linalg.norm(
+        onp.matmul(args, vs) - ws[..., None, :] * vs) < 1e-3))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_ord={}_axis={}_keepdims={}".format(
          jtu.format_shape_dtype_string(shape, dtype), ord, axis, keepdims),
@@ -444,7 +461,6 @@ def testLuGrad(self, shape, dtype, rng):
 
   @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()