Nutpie and PYMC Hurdle Gamma Distribution

[zwelitunyiswa]

Pardon my ignorance, but I cannot get nutpie to sample from PYMC's hurdle-gamma likelihood.

I get "TypeError: The fgraph of ScalarLoop must be exclusively composed of scalar operations."

Is this just not possible at this time?

[aseyboldt]

This looks like an issue with the numba backend of pytensor.
Can you provide an example I can run to further debug this?

In the meantime, you can also give the jax backend a go:

compiled = nutpie.compile_pymc_model(model, backend="jax", gradient_backend="jax")
trace = nutpie.sample(compiled)

[zwelitunyiswa]

Sure. Here is a simple model that replicates the error:

x = pm.HurdleGamma.dist(psi=.2, alpha=1, beta=2)
x_draws = pm.draw(x, draws=500, random_seed=1)
data = pd.DataFrame({'variable':x_draws})

with pm.Model() as model:
    
    alpha = pm.HalfNormal("alpha", sigma=2)
    beta = pm.HalfNormal("beta", sigma=2)
    psi = pm.Beta("psi", alpha=2, beta=2)

    obs = pm.HurdleGamma("obs", psi=psi, alpha=alpha, beta=beta, observed=data)

    idata = pm.sample(nuts_sampler='nutpie')

This fails with the same "TypeError: The fgraph of ScalarLoop must be exclusively composed of scalar operations."

[zwelitunyiswa]

compiled = nutpie.compile_pymc_model(model, backend="jax", gradient_backend="jax")
trace = nutpie.sample(compiled)

I tried that with the simple model above and I got this a "RuntimeError: All initialization points failed" error. Image attached.

[aseyboldt]

Sorry for the delay.
I opened a pytensor issue for the problem.

[zwelitunyiswa]

No worries. Thank you for looking into it and opening the issue.

[zwelitunyiswa]

@aseyboldt I installed the current dev version of pytensor as I saw that they had resolved the issue you posted there #pymc-devs/pytensor#1130.

The hurdle-gamma now runs via Nutpie. I tested it on the sample script ##163 (comment) and it ran. However, it drops down into object mode for same parameters and is slower than pymc's native sampler for this sample code and on some proprietary data and model I have (10 seconds via pymc default sampler vs 12 min for nutpie for my prop).

I am not sure if this should go to another issue and this one should be closed. It seems that it is more a problem with the hurdle-gamma implementation in PYMC #pymc-devs/pymc#7630 (comment) . Please let me know if I should close this issue.

[aseyboldt]

Thanks for the update.
I think you are right, there still seems to be a missing implementation for an op in pytensor in the gradient of the log density. I'll open a separate issue for that.

[ricardoV94]

We should also just implement hurdles properly. There's no reason to involve the cdf as it's just a trick for Mixture to work with a Continuous and the DiracDelta(0) components.

We are truncating the continuous component at epsilon, and that's why those Ops show up in the grad.

[zaxtax]

@ricardoV94 to implement it properly, you mean write the density in a straightforward way instead of using all this Mixture and truncation logic in _hurdle_mixture right?

[ricardoV94]

@ricardoV94 to implement it properly, you mean write the density in a straightforward way instead of using all this Mixture and truncation logic in _hurdle_mixture right?

Yup 👍