Performance improvement in log_det

[frazane]

Type of changes

• Bug fix
• New feature
• Documentation / docstrings
• Tests
• Performance

Checklist

• I've formatted the new code by running poetry run pre-commit run --all-files --show-diff-on-failure before committing.
• I've added tests for new code.
• I've added docstrings for the new code.

Description

Performance improvement in the computation of log_det from LinearOperator.

Edit: unfortunately computing the product before the log can also result in overflow. For small matrices this is not an issue but it is for large ones.

import timeit

import jax
import jax.numpy as jnp
import jax.random as jr
from gpjax.linops import DenseLinearOperator

jax.config.update("jax_enable_x64", True)

key = jr.PRNGKey(123)
cov = jr.normal(key, (100, 100)).astype(jnp.float64)
cov = cov @ cov.T

linop = DenseLinearOperator(cov)
root = linop.to_root()
expr1 = "jnp.sum(jnp.log(root.diagonal()))"
res1 = jnp.sum(jnp.log(root.diagonal()))
time1 = timeit.timeit(expr1, number=1000, globals=globals())

expr2 = "jnp.log(jnp.prod(root.diagonal()))"
res2 = jnp.log(jnp.prod(root.diagonal()))
time2 = timeit.timeit(expr2, number=1000, globals=globals())

print(f"'{expr1}': {time1:.3}s")
print(f"'{expr2}': {time2:.3}s")
assert res1 == res2

'jnp.sum(jnp.log(root.diagonal()))': 0.0363s
'jnp.log(jnp.prod(root.diagonal()))': 0.0178s

[thomaspinder]

Is this ready for merging @frazane ?

[frazane]

@thomaspinder unfortunately I realized there's an issue with numerical stability: floating point overflow occurs for very large matrices (large n), because jnp.prod(root.diagonal()) produces very large numbers.

I would not merge this PR for now. Maybe we could think of some conditional computation based on the size of the matrix. Or do some split-apply-combine kind of thing but I fear this could defeat the purpose of improving performance.