Solver for very stiff Neural Ordinary Differential Equations

I am currently implementing a physically constrained neural ordinary differential equation. The application is reaction kinetics so you already know it's going to be a very stiff system. The governing differential equation is a function of the state Y, the output of two neural networks $M_{1}$ and $M_{2}$ with Y as input (equinox models) and some constant parameters.

$$ \frac{dY}{dt} = f \left(Y, M_{1}(Y), M_{2}(Y), C\right) $$

The two neural networks are pre-trained to produce reasonable outputs. After playing around with atol, rtol, and max_steps of the nonlinear_solver I can solve the equation with the ImplicitEuler method in Diffrax. I can also compare the result y(t) with my target values, compute a loss value, and get the gradient of the loss with respect to the network parameters.

After a few iterations, the diffrax solver throws an error that the implicit method diverges. So I checked if the updates to the neural networks produced some Nans/Infs which was not the case. With the same ode_fun and networks the scipy.integrate.solve_ivp BDF-method works just fine.

When setting the throw=False in diffeqsolve I can see that the solver fails at timesteps with very large gradients or almost zero gradients. Due to the nature of the problem, gradients can span many orders of magnitude.

I have seen that the implementation of the BDF Solver is planed #8 . Maybe that solver is more suitable, but then again I don't have a math background. My question is, should I wait for that implementation, adjust parameters of diffeqsolve further or is there another way around this?

64bit precision is activated via: config.update("jax_enable_x64", True)

I am thankful for any recommendations!

Start with diffeqsolve and the ImplicitEuler solver configuration, including atol, rtol, max_steps, nonlinear_solver, and throw=False. Compare the reported failures at large or nearly zero gradients with the stated scipy.integrate.solve_ivp BDF behavior. Done would require identifying a supported way to solve this stiff neural ODE or documenting that another solver is needed, but the issue provides no reproducer or named test.