Repository: google/jax

Issue:
Feature request for computing eigenvalues of non-hermitian matrix.
I'm basically looking for np.linalg.eig to be implemented in jax or jax.lax.  I personally am looking for a speedup in the computation, and am not interested the derivative at this time.
Thanks!


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
@@ -39,6 +39,10 @@ def cholesky(x, symmetrize_input=True):
     x = symmetrize(x)
   return np.tril(cholesky_p.bind(x))
 
+def eig(x):
+  w, vl, vr = eig_p.bind(x)
+  return w, vl, vr
+
 def eigh(x, lower=True, symmetrize_input=True):
   if symmetrize_input:
     x = symmetrize(x)
@@ -117,6 +121,47 @@ def cholesky_cpu_translation_rule(c, operand):
 
 xla.backend_specific_translations['cpu'][cholesky_p] = cholesky_cpu_translation_rule
 
+# Asymmetric eigendecomposition
+
+def eig_impl(operand):
+  return xla.apply_primitive(eig_p, operand)
+
+def eig_translation_rule(c, operand):
+  raise NotImplementedError(
+    "Nonsymmetric eigendecomposition is only implemented on the CPU backend")
+
+def eig_abstract_eval(operand):
+  if isinstance(operand, ShapedArray):
+    if operand.ndim < 2 or operand.shape[-2] != operand.shape[-1]:
+      raise ValueError("Argument to nonsymmetric eigendecomposition must have "
+                       "shape [..., n, n], got shape {}".format(operand.shape))
+
+    batch_dims = operand.shape[:-2]
+    n = operand.shape[-1]
+    vl = vr = ShapedArray(batch_dims + (n, n), operand.dtype)
+    w = ShapedArray(batch_dims + (n,), lax.lax._complex_basetype(operand.dtype))
+  else:
+    w = vl = vr = operand
+  return core.AbstractTuple((w, vl, vr))
+
+def eig_cpu_translation_rule(c, operand):
+  out = lapack.jax_geev(c, operand)
+  return c.Tuple(c.GetTupleElement(out, 0), c.GetTupleElement(out, 1),
+                 c.GetTupleElement(out, 2))
+
+def eig_batching_rule(batched_args, batch_dims):
+  x, = batched_args
+  bd, = batch_dims
+  x = batching.bdim_at_front(x, bd)
+  return eig_p.bind(x), 0
+
+eig_p = Primitive('eig')
+eig_p.def_impl(eig_impl)
+eig_p.def_abstract_eval(eig_abstract_eval)
+xla.translations[eig_p] = eig_translation_rule
+xla.backend_specific_translations['cpu'][eig_p] = eig_cpu_translation_rule
+batching.primitive_batchers[eig_p] = eig_batching_rule
+
 
 # Symmetric/Hermitian eigendecomposition
 
@@ -132,7 +177,8 @@ def eigh_abstract_eval(operand, lower):
   if isinstance(operand, ShapedArray):
     if operand.ndim < 2 or operand.shape[-2] != operand.shape[-1]:
       raise ValueError(
-        "Argument to symmetric eigendecomposition must have shape [..., n, n]")
+        "Argument to symmetric eigendecomposition must have shape [..., n, n],"
+        "got shape {}".format(operand.shape))
 
     batch_dims = operand.shape[:-2]
     n = operand.shape[-1]
diff --git a/jax/numpy/linalg.py b/jax/numpy/linalg.py
--- a/jax/numpy/linalg.py
+++ b/jax/numpy/linalg.py
@@ -92,6 +92,13 @@ def det(a):
   return sign * np.exp(logdet)
 
 
+@_wraps(onp.linalg.eig)
+def eig(a):
+  a = _promote_arg_dtypes(np.asarray(a))
+  w, vl, vr = lax_linalg.eig(a)
+  return w, vr
+
+
 @_wraps(onp.linalg.eigh)
 def eigh(a, UPLO=None, symmetrize_input=True):
   if UPLO is None or UPLO == "L":


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -104,6 +104,49 @@ def testSlogdet(self, n, dtype, rng):
                             check_dtypes=True, tol=1e-3)
     self._CompileAndCheck(np.linalg.slogdet, args_maker, check_dtypes=True)
 
+  @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 [(0, 0), (4, 4), (5, 5), (50, 50), (2, 6, 6)]
+      for dtype in float_types() + complex_types()
+      for rng in [jtu.rand_default()]))
+  # TODO(phawkins): enable when there is an eigendecomposition implementation
+  # for GPU/TPU.
+  @jtu.skip_on_devices("gpu", "tpu")
+  def testEig(self, shape, dtype, rng):
+    self.skipTest("Test disabled until Jaxlib 0.1.15 is released") # TODO(phawkins)
+    n = shape[-1]
+    args_maker = lambda: [rng(shape, dtype)]
+
+    # Norm, adjusted for dimension and type.
+    def norm(x):
+      norm = onp.linalg.norm(x, axis=(-2, -1))
+      return norm / ((n + 1) * onp.finfo(dtype).eps)
+
+    a, = args_maker()
+    w, v = np.linalg.eig(a)
+    self.assertTrue(onp.all(norm(onp.matmul(a, v) - w[..., None, :] * v) < 100))
+
+    self._CompileAndCheck(partial(np.linalg.eig), args_maker,
+                          check_dtypes=True, rtol=1e-3)
+
+  @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 testEigBatching(self, shape, dtype, rng):
+    self.skipTest("Test disabled until Jaxlib 0.1.15 is released") # TODO(phawkins)
+    shape = (10,) + shape
+    args = rng(shape, dtype)
+    ws, vs = vmap(np.linalg.eig)(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": "_n={}_lower={}".format(
            jtu.format_shape_dtype_string((n,n), dtype), lower),