Incomplete Cholesky

[proteneer]

Does anyone know of a viable incomplete cholesky factorization when the input matrix is not necessarily PSD?

The reference FORTRAN code used by scipy's Newton trust-exact (and other trust radius based methods) is here:

http://www.netlib.org/lapack/explore-html/d1/d7a/group__double_p_ocomputational_ga2f55f604a6003d03b5cd4a0adcfb74d6.html

This is the only function that's been difficult to port over to a jittable jax-implementation, and would enable all sorts of cool things downstream.

[proteneer]

so far best approximation I've found is the projecting down into the nearest pd matrix

def icholesky(H, **kwargs):
    w, v = np.linalg.eigh(H)
    w = np.where(w < 0, 0.0001, w) # make this pd, psd is insufficient
    H_pd = v @ np.eye(3)*w @ v.T

    return jax.scipy.linalg.cholesky(H_pd), np.any(w < 0)

[hawkinsp]

JAX's Cholesky solver calls dpotrf under the hood on CPU. Can you say a bit more about what semantics you would like? You want the behavior where a separate info value says where the decomposition failed? I think we could change the API of the JAX primitive to more closely match LAPACK if it's important.

[proteneer]

Maybe I'm doing something wrong here?

# jax==0.2.6
# jaxlib==0.1.57

import numpy as np
import jax

from scipy.linalg import get_lapack_funcs

H = np.array([[-0.00196721, -0.11516361, -0.5367419 ],
 [-0.11516361,  0.33684297, -0.2710796 ],
 [-0.5367419 , -0.2710796 ,  0.67484898]])

lapack_cholesky, = get_lapack_funcs(('potrf',), (np.array(H),))

print(lapack_cholesky(H)) # no nans, returns tuple
print(jax.scipy.linalg.cholesky(H)) # nans

[proteneer]

(array([[-0.00196721, -0.11516361, -0.5367419 ],
       [ 0.        ,  0.33684297, -0.2710796 ],
       [ 0.        ,  0.        ,  0.67484898]]), 1)
W1201 11:41:02.851007 139898003388224 xla_bridge.py:131] No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[[nan nan nan]
 [ 0. nan nan]
 [ 0.  0. nan]]

[hawkinsp]

Yes, that's intentional: https://github.com/google/jax/blob/bdcf17cdd6be9bef91d2dfbb64e7b8370a6f2295/jax/_src/lax/linalg.py#L326

Now, maybe that was a bad idea and we should have different semantics! We just need to decide what those semantics are. For example, we could change the lax.cholesky primitive to return an info value, but keep jnp.linalg.cholesky with the current all-NaN semantics. (We don't yet have a way to raise an exception from inside a computation, so that's not an option.)

We'd also need to think about how such an info value interacts with things like autodiff.

[proteneer]

Okay the Newton trust-exact method actually needs the non-nan'd U as well as the info:

https://github.com/scipy/scipy/blob/master/scipy/optimize/_trustregion_exact.py#L413

I'm okay with only changing the primitive, and keeping jnp.linalg.cholesky to be consistent with scipy.

[proteneer]

ps - I don't need actually differentiability (since its going to be used in a solver that will be vjp'd via custom_vjp). However, I will need jittability.