Stress test utility for numpyro?

[SamuelBrand1]

Hi everyone,

I'm working on an ODE based model (solver from diffrax) for respiratory viruses fitted on past seasons of data.

At the moment, we occasionally have maximum iteration fails during NUTS warm-up. This is a bit of a classic problem with ODE models, but I was wondering if there was a stress test utility in numpyro which will sample parameters (in the unconstrained domain I guess) and record the parameters that cause a fail in the model log posterior density call (e.g. like this from LogDensityProblems.jl.

I had a look in the docs but I might well have missed something, apols in advance if I have.