Repository: google/jax

Issue:
cholesky not working with vmap on complex arrays
I tried to use `cholesky` on complex arrays with `vmap` but I get `Unimplemented` error.

```
A = np.eye(3).astype(np.complex64)
A_batched = np.stack([A, A])
jax.vmap(np.linalg.cholesky)(A_batched) 

RuntimeError: Unimplemented: Complex types are not implemented in Cholesky; got shape c64[2,3,3]: 
```

This isn't a problem if I don't use vmap.
```
np.linalg.cholesky(A)
DeviceArray([[1.+0.j, 0.+0.j, 0.+0.j],
             [0.+0.j, 1.+0.j, 0.+0.j],
             [0.+0.j, 0.+0.j, 1.+0.j]], dtype=complex64)
```

Is this a bug or am I doing anything wrong? This looks really strange.


Discussion:
Thanks for the question!

What backend are you using (CPU/GPU/TPU)? (I'm guessing CPU.)

It looks like you're hitting [this unimplemented error in XLA](https://github.com/tensorflow/tensorflow/blob/adc3a1b1bec28849f62076e9b4be5c5963e5e5e7/tensorflow/compiler/xla/service/cholesky_expander.cc#L148).

I think the reason it's only coming up when you use `vmap` is that [JAX falls back to an XLA-based implementation](https://github.com/google/jax/blob/29db4203fec6b9796637034d48a62ac84fbcf43f/jax/lax_linalg.py#L146-L149), rather than using a LAPACK-based one as in the un-batched case.

@hawkinsp we could just write a loop in our [`lapack.pyx` bindings](https://github.com/google/jax/blob/29db4203fec6b9796637034d48a62ac84fbcf43f/jaxlib/lapack.pyx#L772) to handle batch dimensions, right? I think that's how we [handle batch dimensions in other kernels](https://github.com/google/jax/blob/29db4203fec6b9796637034d48a62ac84fbcf43f/jaxlib/lapack.pyx#L308).

I'm going to mark this as a "good first issue" because there are plenty of examples in lapack.pyx and lax_linalg.py to pattern-match off of.
(Also I un-assigned myself because I'm not working on this right now.)
will attempt this!

Resolving patch:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -135,20 +135,31 @@ def _nan_like(c, operand):
     nan = c.Constant(onp.array(onp.nan, dtype=dtype))
   return c.Broadcast(nan, shape.dimensions())
 
-def cholesky_cpu_translation_rule(c, operand):
+# TODO(phawkins): remove supports_batching argument after the minimum jaxlib
+# version is 0.1.38.
+def _cholesky_cpu_gpu_translation_rule(potrf_impl, potrf_supports_batching, c,
+                                       operand):
   shape = c.GetShape(operand)
+  batch_dims = shape.dimensions()[:-2]
   dtype = shape.element_type().type
-  if len(shape.dimensions()) == 2 and onp.dtype(dtype) in _cpu_lapack_types:
-    result, info = lapack.potrf(c, operand, lower=True)
-    return c.Select(c.Eq(info, c.ConstantS32Scalar(0)), result,
-                    _nan_like(c, result))
+  if len(batch_dims) == 0 or potrf_supports_batching:
+    result, info = potrf_impl(c, operand, lower=True)
+    ok = c.Eq(info, c.ConstantS32Scalar(0))
+    return _broadcasting_select(c,
+                                c.Reshape(ok, None, batch_dims + (1, 1)), result,
+                                _nan_like(c, result))
   else:
-    # Fall back to the HLO implementation for batched Cholesky decomposition or
-    # unsupported types.
-    # TODO(phawkins): support LAPACK primitives in batched mode.
+    # Fall back to the HLO implementation for batched Cholesky decomposition.
     return c.Cholesky(operand)
 
-xla.backend_specific_translations['cpu'][cholesky_p] = cholesky_cpu_translation_rule
+xla.backend_specific_translations['cpu'][cholesky_p] = partial(
+  _cholesky_cpu_gpu_translation_rule, lapack.potrf,
+  not hasattr(lapack, "jax_potrf"))
+
+# TODO(phawkins): remove after the minimum jaxlib version is 0.1.38.
+if hasattr(cusolver, "potrf"):
+  xla.backend_specific_translations['gpu'][cholesky_p] = partial(
+    _cholesky_cpu_gpu_translation_rule, cusolver.potrf, True)
 
 # Asymmetric eigendecomposition
 
diff --git a/jaxlib/cusolver.py b/jaxlib/cusolver.py
--- a/jaxlib/cusolver.py
+++ b/jaxlib/cusolver.py
@@ -99,6 +99,39 @@ def trsm(c, a, b, left_side=False, lower=False, trans_a=False, conj_a=False,
   return c.GetTupleElement(out, 0)
 
 
+def potrf(c, a, lower):
+  """Cholesky decomposition."""
+  a_shape = c.GetShape(a)
+  dtype = a_shape.element_type()
+  dims = a_shape.dimensions()
+  m, n = dims[-2:]
+  assert m == n
+  batch_dims = tuple(dims[:-2])
+  num_bd = len(batch_dims)
+  batch = _prod(batch_dims)
+
+  lwork, opaque = cusolver_kernels.build_potrf_descriptor(
+      np.dtype(dtype), lower, batch, n)
+  kernel = b"cusolver_potrf"
+
+  out = c.CustomCall(
+      kernel,
+      operands=(a,),
+      shape_with_layout=_Shape.tuple_shape((
+          _Shape.array_shape(
+              dtype, batch_dims + (n, n),
+              (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))),
+          _Shape.array_shape(
+              np.dtype(np.int32), batch_dims, tuple(range(num_bd - 1, -1, -1))),
+          _Shape.array_shape(np.dtype(np.int8), (lwork,), (0,)),
+      )),
+      operand_shapes_with_layout=(_Shape.array_shape(
+          dtype, batch_dims + (n, n),
+          (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))),),
+      opaque=opaque)
+  return c.GetTupleElement(out, 0), c.GetTupleElement(out, 1)
+
+
 def getrf(c, a):
   """LU decomposition."""
   a_shape = c.GetShape(a)


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -72,8 +72,9 @@ def args_maker():
       a = rng(factor_shape, dtype)
       return [onp.matmul(a, np.conj(T(a)))]
 
-    if np.issubdtype(dtype, np.complexfloating) and (
-        len(shape) > 2 or jtu.device_under_test() != "cpu"):
+    if (np.issubdtype(dtype, np.complexfloating) and
+        (jtu.device_under_test() == "tpu" or
+         (jtu.device_under_test() == "cpu" and jax.lib.version < (0, 1, 38)))):
       self.skipTest("Unimplemented case for complex Cholesky decomposition.")
 
     self._CheckAgainstNumpy(onp.linalg.cholesky, np.linalg.cholesky, args_maker,