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.

[fehiepsi]

I think you can draw samples from prior and use the utility log_density to check which of them causes NaN. We don't have a utility for that but please feel free to submit a PR. The feature would be useful for many usage cases.

[SamuelBrand1]

Cheers for the update.