Merge pull request #436 from hawkinsp/master

Move masking of lower triangle of cholesky into Python code.

jax/lax_linalg.py

@@ -37,7 +37,7 @@
 def cholesky(x, symmetrize_input=True):
   if symmetrize_input:
     x = symmetrize(x)
-  return cholesky_p.bind(x)
+  return np.tril(cholesky_p.bind(x))
 
 def eigh(x, lower=True, symmetrize_input=True):
   if symmetrize_input:
@@ -80,7 +80,7 @@ def symmetrize(x): return (x + _H(x)) / 2
 def cholesky_jvp_rule(primals, tangents):
   x, = primals
   sigma_dot, = tangents
-  L = cholesky_p.bind(x)
+  L = np.tril(cholesky_p.bind(x))
 
   # Forward-mode rule from https://arxiv.org/pdf/1602.07527.pdf
   phi = lambda X: np.tril(X) / (1 + np.eye(X.shape[-1], dtype=X.dtype))

jaxlib/lapack.pyx

@@ -333,18 +333,6 @@ cdef void lapack_spotrf(void* out_tuple, void** data) nogil:
 
   spotrf(&uplo, &n, a_out, &n, info)
 
-  # spotrf leaves junk in the part of the triangle that is not written; zero it.
-  cdef int i
-  cdef int j
-  if lower:
-    for i in range(n):
-      for j in range(i):
-        a_out[i * n + j] = 0
-  else:
-    for i in range(n):
-      for j in range(i, n):
-        a_out[i * n + j] = 0
-
 register_cpu_custom_call_target(b"lapack_spotrf", <void*>(lapack_spotrf))
 
 
@@ -362,18 +350,6 @@ cdef void lapack_dpotrf(void* out_tuple, void** data) nogil:
 
   dpotrf(&uplo, &n, a_out, &n, info)
 
-  # dpotrf leaves junk in the part of the triangle that is not written; zero it.
-  cdef int i
-  cdef int j
-  if lower:
-    for i in range(n):
-      for j in range(i):
-        a_out[i * n + j] = 0
-  else:
-    for i in range(n):
-      for j in range(i, n):
-        a_out[i * n + j] = 0
-
 register_cpu_custom_call_target(b"lapack_dpotrf", <void*>(lapack_dpotrf))
 
 
@@ -391,18 +367,6 @@ cdef void lapack_cpotrf(void* out_tuple, void** data) nogil:
 
   cpotrf(&uplo, &n, a_out, &n, info)
 
-  # cpotrf leaves junk in the part of the triangle that is not written; zero it.
-  cdef int i
-  cdef int j
-  if lower:
-    for i in range(n):
-      for j in range(i):
-        a_out[i * n + j] = 0
-  else:
-    for i in range(n):
-      for j in range(i, n):
-        a_out[i * n + j] = 0
-
 register_cpu_custom_call_target(b"lapack_cpotrf", <void*>(lapack_cpotrf))
 
 cdef void lapack_zpotrf(void* out_tuple, void** data) nogil:
@@ -419,18 +383,6 @@ cdef void lapack_zpotrf(void* out_tuple, void** data) nogil:
 
   zpotrf(&uplo, &n, a_out, &n, info)
 
-  # zpotrf leaves junk in the part of the triangle that is not written; zero it.
-  cdef int i
-  cdef int j
-  if lower:
-    for i in range(n):
-      for j in range(i):
-        a_out[i * n + j] = 0
-  else:
-    for i in range(n):
-      for j in range(i, n):
-        a_out[i * n + j] = 0
-
 register_cpu_custom_call_target(b"lapack_zpotrf", <void*>(lapack_zpotrf))
 
 def jax_potrf(c, a, lower=False):