Fix check for missing rv in factorized_joint_logprob

[larryshamalama]

Currently, random variables that are not provided with their value variable equivalent can result in a KeyError in factorized_joint_logprob when given warn_missing_rvs=False. This error rose when working with a graph corresponding to a Gaussian random walk.

import aesara.tensor as at
from aeppl import factorized_joint_logprob

mu = at.random.normal(loc=2, scale=0.1)
sigma = at.random.gamma(shape=1, rate=1)
innovation_dist = at.random.normal(mu[..., None], sigma[..., None], size=(10,))
rv_out = at.cumsum(innovation_dist, axis=-1) # gaussian random walk

The following code chunk yields a KeyError:

factorized_joint_logprob(
    {
        rv_out: rv_out.copy(),
    }, 
    extra_rewrites=TransformValuesOpt({}),
    use_jacobian=True,
    warn_missing_rvs=False
)

but setting warn_missing_rvs=True only issues warnings.

In this PR, so far, I fix the case where users decide to not issue warnings.

[ricardoV94]

Can you add a test? Your example should be fine just assert that the warning raised in one case but not the other, but the output is the same.

[larryshamalama]

Done! There was already a test for this, I just added a verification for warn_missing_rvs=False. Here, is it better to have both commits as one?

[larryshamalama]

Failing test as described in #148:

https://github.com/aesara-devs/aeppl/blob/cc78f30b5ed89e5b247b88d697467a76cc1e424e/tests/test_scan.py#L349

[ricardoV94]

You can just cherry pick (or copy) this commit: https://github.com/aesara-devs/aeppl/commit/e0f9623fa6734061e05850eb14c72e18ddd7027f

[larryshamalama]

You can just cherry pick (or copy) this commit: e0f9623

Done!

[ricardoV94]

You should have the two commits separate.

One for the xfail and the other for the bugfix and new test condition.

[larryshamalama]

You should have the two commits separate.

One for the xfail and the other for the bugfix and new test condition.

Done! Thanks for the clarification

[ricardoV94]

Small nitpick: the test and respective bugfix make sense to go in the same commit

[larryshamalama]

Small nitpick: the test and respective bugfix make sense to go in the same commit

After several attempts and many force pushes, I think that I got it right (?) 😅

[codecov[bot]]

Codecov Report

Merging #151 (832e329) into main (cc78f30) will increase coverage by 0.00%. The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #151   +/-   ##
=======================================
  Coverage   94.88%   94.88%           
=======================================
  Files          12       12           
  Lines        1779     1780    +1     
  Branches      262      263    +1     
=======================================
+ Hits         1688     1689    +1     
  Misses         51       51           
  Partials       40       40           

Impacted Files	Coverage Δ
aeppl/joint_logprob.py	96.92% <100.00%> (+0.04%)	:arrow_up:

---

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update cc78f30...832e329. Read the comment docs.