Repository: pymc-devs/pymc

Issue:
Invalid logp expression if args aren't explicit in DensityDist
## Description of your problem

On the latest pymc version (777622aea4e5edbbc6051c1d1a0d32456e878b59) I get incorret logp graphs that contain sampled random variables if I use model variables in a closure of a density dist:

```python
import numpy as np
import pymc as pm

with pm.Model() as model:
    a = pm.Normal("a")
    pm.DensityDist("b", logp=lambda x: (x - a) ** 2, observed=np.array(3.))
```
The logp function of the model is not deterministic now, because it uses a *sampled* version of `a` in the logp function of `b` instead of the value from the value variable. You can see this in the aesara graph (the normal_rv op should not be in here):

```python
aesara.dprint(model.logpt)
```

```
Elemwise{add,no_inplace} [id A] '__logp'   
 |Elemwise{add,no_inplace} [id B] ''   
 | |Sum{acc_dtype=float64} [id C] ''   
 | | |TensorConstant{[]} [id D]
 | |Sum{acc_dtype=float64} [id E] ''   
 |   |Elemwise{mul,no_inplace} [id F] ''   
 |     |Assert{msg='sigma > 0'} [id G] 'a_logprob'   
 |     | |Elemwise{sub,no_inplace} [id H] ''   
 |     | | |Elemwise{sub,no_inplace} [id I] ''   
 |     | | | |Elemwise{mul,no_inplace} [id J] ''   
 |     | | | | |TensorConstant{-0.5} [id K]
 |     | | | | |Elemwise{pow,no_inplace} [id L] ''   
 |     | | | |   |Elemwise{true_div,no_inplace} [id M] ''   
 |     | | | |   | |Elemwise{sub,no_inplace} [id N] ''   
 |     | | | |   | | |a [id O]
 |     | | | |   | | |TensorConstant{0} [id P]
 |     | | | |   | |TensorConstant{1.0} [id Q]
 |     | | | |   |TensorConstant{2} [id R]
 |     | | | |Elemwise{log,no_inplace} [id S] ''   
 |     | | |   |TensorConstant{2.5066282746310002} [id T]
 |     | | |Elemwise{log,no_inplace} [id U] ''   
 |     | |   |TensorConstant{1.0} [id Q]
 |     | |All [id V] ''   
 |     |   |Elemwise{gt,no_inplace} [id W] ''   
 |     |     |TensorConstant{1.0} [id Q]
 |     |     |TensorConstant{0.0} [id X]
 |     |TensorConstant{1.0} [id Y]
 |Sum{acc_dtype=float64} [id Z] ''   
   |Elemwise{mul,no_inplace} [id BA] ''   
     |Elemwise{pow,no_inplace} [id BB] 'b_logprob'   
     | |Elemwise{sub,no_inplace} [id BC] ''   
     | | |TensorConstant{3.0} [id BD]
     | | |normal_rv{0, (0, 0), floatX, False}.1 [id BE] 'a'   
     | |   |RandomStateSharedVariable(<RandomState(MT19937) at 0x7F660912BC40>) [id BF]
     | |   |TensorConstant{[]} [id BG]
     | |   |TensorConstant{11} [id BH]
     | |   |TensorConstant{0} [id BI]
     | |   |TensorConstant{1.0} [id BJ]
     | |TensorConstant{2} [id BK]
     |TensorConstant{1.0} [id BL]
```

This can be fixed in user code by explicitly letting the DensityDist know about the parameter:
```python
with pm.Model() as model2:
    a = pm.Normal("a")
    pm.DensityDist("b", a, logp=lambda x, a_val: (x - a_val) ** 2, observed=np.array(3.))
```

Finding a bug like this in an actual model took @ferrine and me a couple of hours, so if we could either change `pm.logp` so that this just works or so that we get an error if there are remaining rv ops in a logp graph that would help a lot.

(cc @ricardoV94)

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/pymc-devs/pymc/pull/5614?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
> Merging [#5614](https://codecov.io/gh/pymc-devs/pymc/pull/5614?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (9b7a51d) into [main](https://codecov.io/gh/pymc-devs/pymc/commit/144b0baa1a560cfe1d75ace2a0b8fb70f8a26448?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (144b0ba) will **decrease** coverage by `40.14%`.
> The diff coverage is `80.00%`.

[![Impacted file tree graph](https://codecov.io/gh/pymc-devs/pymc/pull/5614/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://codecov.io/gh/pymc-devs/pymc/pull/5614?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@             Coverage Diff             @@
##             main    #5614       +/-   ##
===========================================
- Coverage   87.68%   47.54%   -40.15%     
===========================================
  Files          76       76               
  Lines       13702    13702               
===========================================
- Hits        12014     6514     -5500     
- Misses       1688     7188     +5500     
```


| [Impacted Files](https://codecov.io/gh/pymc-devs/pymc/pull/5614?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/distributions/logprob.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL2xvZ3Byb2IucHk=) | `62.28% <80.00%> (-34.05%)` | :arrow_down: |
| [pymc/ode/utils.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9vZGUvdXRpbHMucHk=) | `15.62% <0.00%> (-84.38%)` | :arrow_down: |
| [pymc/step\_methods/mlda.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvbWxkYS5weQ==) | `12.53% <0.00%> (-83.85%)` | :arrow_down: |
| [pymc/variational/updates.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC91cGRhdGVzLnB5) | `11.33% <0.00%> (-80.79%)` | :arrow_down: |
| [pymc/smc/smc.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zbWMvc21jLnB5) | `18.62% <0.00%> (-79.37%)` | :arrow_down: |
| [pymc/gp/gp.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9ncC9ncC5weQ==) | `18.68% <0.00%> (-74.50%)` | :arrow_down: |
| [pymc/tuning/starting.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy90dW5pbmcvc3RhcnRpbmcucHk=) | `19.51% <0.00%> (-73.18%)` | :arrow_down: |
| [pymc/step\_methods/slicer.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvc2xpY2VyLnB5) | `26.76% <0.00%> (-69.02%)` | :arrow_down: |
| [pymc/gp/cov.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9ncC9jb3YucHk=) | `29.20% <0.00%> (-68.88%)` | :arrow_down: |
| [pymc/variational/inference.py](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC9pbmZlcmVuY2UucHk=) | `19.38% <0.00%> (-67.86%)` | :arrow_down: |
| ... and [48 more](https://codecov.io/gh/pymc-devs/pymc/pull/5614/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | |

ricardoV94: Failing test seems unrelated

Patch:
diff --git a/pymc/distributions/logprob.py b/pymc/distributions/logprob.py
index 1b2846845..bb76f444e 100644
--- a/pymc/distributions/logprob.py
+++ b/pymc/distributions/logprob.py
@@ -24,6 +24,7 @@ from aeppl.logprob import logcdf as logcdf_aeppl
 from aeppl.logprob import logprob as logp_aeppl
 from aeppl.transforms import TransformValuesOpt
 from aesara.graph.basic import graph_inputs, io_toposort
+from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.subtensor import (
     AdvancedIncSubtensor,
     AdvancedIncSubtensor1,
@@ -223,6 +224,26 @@ def joint_logpt(
         tmp_rvs_to_values, extra_rewrites=transform_opt, use_jacobian=jacobian, **kwargs
     )
 
+    # Raise if there are unexpected RandomVariables in the logp graph
+    # Only SimulatorRVs are allowed
+    from pymc.distributions.simulator import SimulatorRV
+
+    unexpected_rv_nodes = [
+        node
+        for node in aesara.graph.ancestors(list(temp_logp_var_dict.values()))
+        if (
+            node.owner
+            and isinstance(node.owner.op, RandomVariable)
+            and not isinstance(node.owner.op, SimulatorRV)
+        )
+    ]
+    if unexpected_rv_nodes:
+        raise ValueError(
+            f"Random variables detected in the logp graph: {unexpected_rv_nodes}.\n"
+            "This can happen when DensityDist logp or Interval transform functions "
+            "reference nonlocal variables."
+        )
+
     # aeppl returns the logpt for every single value term we provided to it. This includes
     # the extra values we plugged in above, so we filter those we actually wanted in the
     # same order they were given in.


Test patch:
diff --git a/pymc/tests/test_logprob.py b/pymc/tests/test_logprob.py
index 1afd513a0..f6ff39a2c 100644
--- a/pymc/tests/test_logprob.py
+++ b/pymc/tests/test_logprob.py
@@ -28,6 +28,7 @@ from aesara.tensor.subtensor import (
     Subtensor,
 )
 
+from pymc import DensityDist
 from pymc.aesaraf import floatX, walk_model
 from pymc.distributions.continuous import HalfFlat, Normal, TruncatedNormal, Uniform
 from pymc.distributions.discrete import Bernoulli
@@ -217,3 +218,12 @@ def test_model_unchanged_logprob_access():
     model.logpt()
     new_inputs = set(aesara.graph.graph_inputs([c]))
     assert original_inputs == new_inputs
+
+
+def test_unexpected_rvs():
+    with Model() as model:
+        x = Normal("x")
+        y = DensityDist("y", logp=lambda *args: x)
+
+    with pytest.raises(ValueError, match="^Random variables detected in the logp graph"):
+        model.logpt()
diff --git a/pymc/tests/test_parallel_sampling.py b/pymc/tests/test_parallel_sampling.py
index 78dad2c68..d321ef868 100644
--- a/pymc/tests/test_parallel_sampling.py
+++ b/pymc/tests/test_parallel_sampling.py
@@ -201,11 +201,11 @@ def test_spawn_densitydist_bound_method():
     N = 100
     with pm.Model() as model:
         mu = pm.Normal("mu", 0, 1)
-        normal_dist = pm.Normal.dist(mu, 1, size=N)
 
-        def logp(x):
+        def logp(x, mu):
+            normal_dist = pm.Normal.dist(mu, 1, size=N)
             out = pm.logp(normal_dist, x)
             return out
 
-        obs = pm.DensityDist("density_dist", logp=logp, observed=np.random.randn(N), size=N)
+        obs = pm.DensityDist("density_dist", mu, logp=logp, observed=np.random.randn(N), size=N)
         pm.sample(draws=10, tune=10, step=pm.Metropolis(), cores=2, mp_ctx="spawn")