[JAX] Cholesky decomposition error 230318 기준

현재 기준(230318)으로 cholesky decomposition을 사용할 때 행렬 크기가 50정도 넘어가면 nan을 출력하는 오류가 있다.

 

jnp.linalg.cholesky(K)
jax.random.multivariate_normal(subkeys[0], np.zeros((N_samples,)), K)

 

이 때문에 cholesky decomposition을 쓰는 다른 함수들도 영향을 받았는데

jax.random.multivariate_normal의 경우에도 랜덤하게 추출하는 과정에서 cholesky decomposition을 쓴다.

cholesky decomposition은 어쩔 수 없을 것 같고

 

jax.random.multivariate_normal(subkeys[0], np.zeros((N_samples,)), K, method='svd')

jax.random.multivariate_normal의 경우에는 다른 메소드(위에서는 svd)를 쓰면 에러없이 사용할 수 있다.