Fix NaN gradients of the CARMA kernel - #284

Fixes #228.

CARMA.__init__ builds the observation model from square roots that are computed for every root, real or complex, and then picks the right one with jnp.where. For a real root the complex terms are sqrt(0), and for some complex cases (e.g. CARMA(2,0)) the h1 radicand is exactly 0 as well. The derivative of jnp.sqrt at 0 is infinite, and jnp.where multiplies it by 0 for the unused branch, which gives NaN. So jax.grad of a log likelihood with a CARMA kernel was NaN for CARMA(1,0), CARMA(2,0), CARMA(2,1) and CARMA(3,1) alike, which is what the issue hit in NumPyro.

The fix adds a small _safe_sqrt that returns 0 with a zero derivative for non-positive input, and feeds it 0 for the terms that jnp.where discards. The values that are used do not change. As a side effect, a radicand that comes out as a tiny negative number from rounding now gives 0 instead of a NaN in the forward pass.

With the example from the issue, the gradient now matches central finite differences:

alpha: -34.0916 (grad) vs -34.0916 (finite diff)
beta: -25688.8092 vs -25688.8092

Tests: test_carma_grad in tests/test_kernels/test_quasisep.py checks that the gradient is finite and passes jax.test_util.check_grads for a real root, two complex cases and a CARMA(3,1) with both kinds of roots. All four fail on main. tests/test_kernels/test_quasisep.py passes locally, black and ruff at the pre-commit versions are clean, and there is a news fragment in news/228.bugfix.