Building Non-Standard Regression Models with JAXUsing JAX’s automatic differentiation to train a Zero-Inflated Generalised Poisson regression model