Log-pseudo-determinant for matrix times a constant

[Seb-Lorek]

I have a question regarding the calculation of a pseudo determinant using jax.numpy see my question on stack exchange for some math background on the question. When using numpy the property det(cK)_{+}=c^{rank(K)}det(K)_{+} (or the log-transformation of this) holds. But when using jax.numpy there is a quite significant numerical difference. In numpy

import numpy as np
scale = np.array([0.5])
diff_mat = np.diff(np.eye(20), n=2, axis=0)
pen = np.dot(diff_mat.T, diff_mat)
prec = pen / (scale**2)

# based on pen
eigval_pen = np.linalg.eigvalsh(pen)
mask_pen = eigval_pen > 1e-6
rank_pen = np.linalg.matrix_rank(pen)
selec_pen = np.where(mask_pen, eigval_pen, 1.0)
log_pdet_pen1 =  np.sum(np.log(selec_pen), axis=-1)
log_pdet_prec1 = log_pdet_pen1 - rank_pen * 2 * np.log(scale)

# based on prec 
eigval_prec = np.linalg.eigvalsh(prec)
mask_prec = eigval_prec > 1e-6
rank_prec = np.linalg.matrix_rank(prec)
selec_prec = np.where(mask_prec, eigval_prec, 1.0)
log_pdet_prec2 = np.sum(np.log(selec_prec), axis=-1)

np.isclose(log_pdet_prec1, log_pdet_prec2, rtol=1e-04)

results in

array([ True]).

But in jax.numpy

import jax.numpy as jnp
scale = jnp.array([0.5])
diff_mat = jnp.diff(jnp.eye(20), n=2, axis=0)
pen = jnp.dot(diff_mat.T, diff_mat)
prec = pen / (scale**2)

# based on penalty 
eigval_pen = jax.numpy.linalg.eigvalsh(pen)
mask_pen = eigval_pen > 1e-6
rank_pen = jnp.linalg.matrix_rank(pen)
selec_pen = jnp.where(mask_pen, eigval_pen, 1.0)
log_pdet_pen1 =  jnp.sum(jnp.log(selec_pen), axis=-1)
log_pdet_prec1 = log_pdet_pen1 - rank_pen * 2 * jnp.log(scale)

# based on prec 
eigval_prec = jax.numpy.linalg.eigvalsh(prec)
mask_prec = eigval_prec > 1e-6
rank_prec = jnp.linalg.matrix_rank(prec)
selec_prec = jnp.where(mask_prec, eigval_prec, 1.0)
log_pdet_prec2 = jnp.sum(jnp.log(selec_prec), axis=-1)

jnp.isclose(log_pdet_prec1, log_pdet_prec2, rtol=1e-04)

results in

Array([False], dtype=bool).

And the difference is quite significant.

log_pdet_prec1 - log_pdet_prec2

results in

Array([12.973423], dtype=float32).

Any ideas where this difference stems from ? Due to float32 ?

[jakevdp]

What backend are you running this on, and what version of JAX are you using? I tried running your code on both a Colab CPU and Colab T4 GPU runtime, with X64 both True and False, and in all cases the JAX code results in [True]

[Seb-Lorek]

Running this on local CPU and the jax version is 0.4.12.

[jakevdp]

I tried running your code with JAX 0.4.12 on a CPU, and the output both on Colab and my OSX laptop is Array([ True], dtype=bool).

Given that, I suspect the difference may have to do with the linear algebra libraries (BLAS/LAPACK) that you have installed. On a CPU, JAX uses the same lapack libraries as scipy. Which scipy version do you have installed? And do you have any particular lapack version installed (e.g. if you're using anaconda, you probably have lapack from MKL)?

[Seb-Lorek]

Scipy version 1.10.1. Yes I'm using anaconda, however when comparing my output from np.__config__.show()with resource my output diverges and results in:

blas_info:
    libraries = ['cblas', 'blas', 'cblas', 'blas']
    library_dirs = ['/Users/SebLorek/anaconda3/envs/thesis/lib']
    include_dirs = ['/Users/SebLorek/anaconda3/envs/thesis/include']
    language = c
    define_macros = [('HAVE_CBLAS', None)]
blas_opt_info:
    define_macros = [('NO_ATLAS_INFO', 1), ('HAVE_CBLAS', None)]
    libraries = ['cblas', 'blas', 'cblas', 'blas']
    library_dirs = ['/Users/SebLorek/anaconda3/envs/thesis/lib']
    include_dirs = ['/Users/SebLorek/anaconda3/envs/thesis/include']
    language = c
lapack_info:
    libraries = ['lapack', 'blas', 'lapack', 'blas']
    library_dirs = ['/Users/SebLorek/anaconda3/envs/thesis/lib']
    language = f77
lapack_opt_info:
    libraries = ['lapack', 'blas', 'lapack', 'blas', 'cblas', 'blas', 'cblas', 'blas']
    library_dirs = ['/Users/SebLorek/anaconda3/envs/thesis/lib']
    language = c
    define_macros = [('NO_ATLAS_INFO', 1), ('HAVE_CBLAS', None)]
    include_dirs = ['/Users/SebLorek/anaconda3/envs/thesis/include']
Supported SIMD extensions in this NumPy install:
    baseline = SSE,SSE2,SSE3
    found = SSSE3,SSE41,POPCNT,SSE42,AVX,F16C,FMA3,AVX2
    not found = AVX512F,AVX512CD,AVX512_KNL,AVX512_SKX,AVX512_CLX,AVX512_CNL,AVX512_ICL

[Seb-Lorek]

I essentially just set up a new virtual environment and reinstalled all dependencies. And now the code evaluates in both cases to Array([ True], dtype=bool). Moreover np.__config__.show() now results in

openblas64__info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
blas_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None)]
    runtime_library_dirs = ['/usr/local/lib']
openblas64__lapack_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
lapack_ilp64_opt_info:
    libraries = ['openblas64_', 'openblas64_']
    library_dirs = ['/usr/local/lib']
    language = c
    define_macros = [('HAVE_CBLAS', None), ('BLAS_SYMBOL_SUFFIX', '64_'), ('HAVE_BLAS_ILP64', None), ('HAVE_LAPACKE', None)]
    runtime_library_dirs = ['/usr/local/lib']
Supported SIMD extensions in this NumPy install:
    baseline = SSE,SSE2,SSE3
    found = SSSE3,SSE41,POPCNT,SSE42,AVX,F16C,FMA3,AVX2
    not found = AVX512F,AVX512CD,AVX512_KNL,AVX512_SKX,AVX512_CLX,AVX512_CNL,AVX512_ICL

So the issue was related to the linear algebra libraries as you said. Thanks so much for your help, if you have some more input I'm happy to hear it.