Allow Minibatch logp on derived RVs

Author: ricardoV94

pymc/variational/minibatch_rv.py

@@ -20,7 +20,8 @@
 from pytensor.graph import Apply, Op
 from pytensor.tensor import NoneConst, TensorVariable, as_tensor_variable
 
-from pymc.logprob.abstract import MeasurableOp, _logprob, _logprob_helper
+from pymc.logprob.abstract import MeasurableOp, _logprob
+from pymc.logprob.basic import logp
 
 
 class MinibatchRandomVariable(MeasurableOp, Op):
@@ -99,4 +100,4 @@ def get_scaling(total_size: Sequence[Variable], shape: TensorVariable) -> Tensor
 def minibatch_rv_logprob(op, values, *inputs, **kwargs):
     [value] = values
     rv, *total_size = inputs
-    return _logprob_helper(rv, value, **kwargs) * get_scaling(total_size, value.shape)
+    return logp(rv, value, **kwargs) * get_scaling(total_size, value.shape)

tests/variational/test_minibatch_rv.py

@@ -13,6 +13,7 @@
 #   limitations under the License.
 import numpy as np
 import pytensor
+import pytensor.tensor as pt
 import pytest
 
 from scipy import stats as st
@@ -186,3 +187,12 @@ def test_minibatch_parameter_and_value(self):
         with m:
             pm.set_data({"AD": rng.normal(size=1000)})
         assert logp_fn(ip) != logp_fn(ip)
+
+    def test_derived_rv(self):
+        """Test we can obtain a minibatch logp out of a derived RV."""
+        dist = pt.clip(pm.Normal.dist(0, 1, size=(1,)), -1, 1)
+        mb_dist = create_minibatch_rv(dist, total_size=(2,))
+        np.testing.assert_allclose(
+            pm.logp(mb_dist, -1).eval(),
+            pm.logp(dist, -1).eval() * 2,
+        )