refactor: move _precition_mv_normal_logp into a seperate function

[Michal-Novomestsky]

Addresses #7894 as part of pymc-devs/pymc-extras#532

---

📚 Documentation preview 📚: https://pymc--7895.org.readthedocs.build/en/7895/

[welcome]

]
💖 Thanks for opening this pull request! 💖 The PyMC community really appreciates your time and effort to contribute to the project. Please make sure you have read our Contributing Guidelines and filled in our pull request template to the best of your ability.

[ricardoV94]

Let's duplicate it in pymc-extras first to see if we really need it refactored. In theory you should be able to introduce a PrecisionMVNormal RV and call logp on it to get this implementation. Anyway it will be more obvious if we need the refactor once the dependent work on pymc-extras is ready to be merged.

[Michal-Novomestsky]

@ricardoV94 Thoughts on this having seen the main PR now? Also I'm a little unsure about creating a PrecisionalMVNormal RV inside the logp function - I tried doing that with MvNormal(tau=...), but then it complained that the RV already exists (I believe it happens because the logp gets called repeatedly and ig the garbage collector doesn't wipe the RV in time. Please correct me if I'm wrong, but here I feel that simply calling precision_mv_normal_logp is much more lightweight, especially given speed is a concern.

[Michal-Novomestsky]

pre-commit.ci autofix

[Michal-Novomestsky]

@ricardoV94 Thoughts on this having seen the main PR now? Also I'm a little unsure about creating a PrecisionalMVNormal RV inside the logp function - I tried doing that with MvNormal(tau=...), but then it complained that the RV already exists (I believe it happens because the logp gets called repeatedly and ig the garbage collector doesn't wipe the RV in time. Please correct me if I'm wrong, but here I feel that simply calling precision_mv_normal_logp is much more lightweight, especially given speed is a concern.

@ricardoV94 Just making sure you're ok with the above?

[ricardoV94]

What I'm saying is that you can already get the logp function wherever it is implemented by following the logp API which is:

import pymc as pm
import pytensor.tensor as pt
from pymc.distributions.multivariate import PrecisionMvNormalRV

# There's a bug here, it should be possible to do
# rv = PrecisionMvNormalRV.rv_op(mean=pt.ones(3), tau=pt.eye(3))
# but we wrote the signature wrong

_, rv = PrecisionMvNormalRV.rv_op(mean=pt.ones(3), tau=pt.eye(3)).owner.outputs
logp = pm.logp(rv, pt.zeros(3))
logp.eval()

This is the API to access the logp function. It should work just fine for you inside a logp expression. We do it here as well:

pymc/pymc/distributions/discrete.py

Line 734 in 0960323

return pt.switch(pt.gt(alpha, 1e10), logp(Poisson.dist(mu=mu), value), negbinom)

[ricardoV94]

In either case this PR is fine, I just want to do it after the work on the pymc-extras is ready. Until then just copy the logp function there

[Michal-Novomestsky]

In either case this PR is fine, I just want to do it after the work on the pymc-extras is ready. Until then just copy the logp function there

Sounds good - should I just raise an issue for now in that case?

[ricardoV94]

Sounds good - should I just raise an issue for now in that case?

Just add the code in the same PR where you need it

[Michal-Novomestsky]

Just add the code in the same PR where you need it

Oh yes I see, just done now.