Inverse Gaussian distribution

Hi everyone,
I am trying to implement a model where I have to use an Inverse Gaussian distribution. As far as I could see, Inverse Gaussian is available in Scipy but it is using numpy methods (it is not available in JAX version of Scipy) and when I wrap it inside a Distribution class in numpyro it does not work as it is not a JAX method and numpyro can not trace it. Is there any way to make this work and use Inverse Gaussian distribution inside numpyro?


Hi @ideznaby, after this PR, you will be able to use a bunch of TFP distributions in NumPyro, including InverseGaussian. :slight_smile: