Explosion of steps for specific parameter values

I have been experiencing odd integration failures in large sets of solves of relatively small simply systems of equations. I have narrowed this down to a small example:

import diffrax
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np

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

a = jnp.array(
    [
        6.026932645397832,
        4.41195014234956,
        5.884199824299863,
        3.673504195449191,
        4.17957753821087,
    ]
)
b = jnp.float64(
    -2.823760940491063
)


def xdot(t, x, _):
    c = jnp.exp(a[1]) - x[0]
    d = x[1] / (c + jnp.exp(b))
    dxdt = jnp.asarray(
        (
            jnp.exp(a[3]) * d * c - jnp.exp(a[4]) * x[0],
            jnp.exp(a[3]) * (jnp.exp(a[0]) - x[1]),
        )
    )
    return dxdt


def xdot_lse(t, x, _):
    c = jnp.log(jnp.exp(a[1]) - x[0])
    d = jnp.log(x[1]) - jax.nn.logsumexp(jnp.array([b, c]))
    dxdt = jnp.asarray(
        (
            jnp.exp(a[3] + d + c) - jnp.exp(a[4]) * x[0],
            jnp.exp(a[3]) * (jnp.exp(a[0]) - x[1]),
        )
    )
    return dxdt


y0 = jnp.array([4.1154706432848185, 6.831774897154676])
ts = np.concatenate((np.array([0]), np.logspace(-6, 2, 20)))

sol = diffrax.diffeqsolve(
    diffrax.ODETerm(xdot_lse),
    solver=diffrax.Kvaerno5(),
    t0=0.0,
    t1=ts[-1],
    dt0=1e-8,
    y0=y0,
    stepsize_controller=diffrax.PIDController(
        atol=1e-8,
        rtol=1e-6,
        pcoeff=0.4,
        icoeff=0.3,
        dcoeff=0,
    ),
    max_steps=int(1e5),
    saveat=diffrax.SaveAt(ts=ts),
    throw=False,
)
x = sol.ys[-2, :]

for mag in np.logspace(-14, -4, 6):
    xx = np.linspace(-mag, mag, 200)
    fx_lse = jnp.asarray(
        [xdot_lse(0.0, x + jnp.asarray([eps, 0.0]), None)[0] for eps in xx]
    )
    fx = jnp.asarray(
        [xdot(0.0, x + jnp.asarray([eps, 0.0]), None)[0] for eps in xx]
    )
    f, axes = plt.subplots(1, 2)
    axes[0].plot(xx, fx_lse, marker='o', label='xdot_lse')
    axes[0].plot(xx, fx, marker='.', label='xdot')
    axes[0].set_ylabel('value of dx/dt[0] at x+eps')
    axes[0].set_xlabel('eps')
    axes[0].legend()
    axes[1].plot(xx, fx - fx_lse, marker='o')
    axes[1].set_ylabel('implementation difference dx/dt[0] at x+eps')
    axes[1].set_xlabel('eps')
    plt.tight_layout()
    plt.show()

The example should fail at the last time-point with about ~50k rejected steps and ~50k accepted steps. Minuscule changes to the parameters, e.g. changing the first entry in a from 6.026932645397832 to 6.02693264539783 allows the system to be solved in ~90 steps. This is odd as the systems is pretty close to a steady state when the integration fails and should be easy to integrate.

I initially thought this might be the result of some numerical instability, but I'm no longer convinced that this is the case. For example, changing d to d = x[1]/(c + jnp.exp(b[0])) (implemented in xdot) resolves the integration failure, but doesn't result in any appreciably improvement in numerical stability with which the right hand side can be evaluated (see plots generated at the end of the script). The magnitude of changes that I see are in the range of 1e-11 to 1e-12, which in my understanding shouldn't matter too much for the tolerances that I am using. Therefore, my conclusion is that I might be hitting some weird numerical edge-case.

Run the supplied reproducer and compare the step counts and RHS behavior for xdot and xdot_lse across the parameter values shown. Trace the solver's rejected and accepted steps around the final time point; done means the parameter-sensitive step explosion is explained and corrected without breaking the reported integration behavior.