keras/src/backend/jax/linalg.py/cholesky
def cholesky(a):
    out = jnp.linalg.cholesky(a)
    try:
        # In eager mode, raise for nan to
        # achieve behavior consistency with numpy
        if jnp.any(jnp.isnan(out)):
            raise ValueError(
                "Cholesky decomposition failed. "
                "The input might not be a valid "
                "positive definite matrix."
            )
    except jax.errors.TracerBoolConversionError:
        # Cannot raise for nan in tracing mode
        pass
    return out