You might want to add some print statements to your model to see why mu is invalid for Poisson likelihood. It seems to get inf or nan. I guess jnp.exp(beta * df.metastasized.values) somehow blow up? You might want to use bounded_exp here.
You might want to add some print statements to your model to see why mu is invalid for Poisson likelihood. It seems to get inf or nan. I guess jnp.exp(beta * df.metastasized.values) somehow blow up? You might want to use bounded_exp here.