Fallback to Aeppl when logp is not implemented

Author: ricardoV94

pymc/distributions/logprob.py

@@ -269,11 +269,23 @@ def joint_logpt(
     return logp_var
 
 
-def logp(rv, value):
+def logp(rv: TensorVariable, value) -> TensorVariable:
     """Return the log-probability graph of a Random Variable"""
 
     value = at.as_tensor_variable(value, dtype=rv.dtype)
-    return logp_aeppl(rv, value)
+    try:
+        return logp_aeppl(rv, value)
+    except NotImplementedError:
+        try:
+            value = rv.type.filter_variable(value)
+        except TypeError as exc:
+            raise TypeError(
+                "When RV is not a pure distribution, value variable must have the same type"
+            ) from exc
+        try:
+            return factorized_joint_logprob({rv: value}, warn_missing_rvs=False)[value]
+        except Exception as exc:
+            raise NotImplementedError("PyMC could not infer logp of input variable.") from exc
 
 
 def logcdf(rv, value):

pymc/tests/test_logprob.py

@@ -31,7 +31,13 @@
 
 from pymc import DensityDist
 from pymc.aesaraf import floatX, walk_model
-from pymc.distributions.continuous import HalfFlat, Normal, TruncatedNormal, Uniform
+from pymc.distributions.continuous import (
+    HalfFlat,
+    LogNormal,
+    Normal,
+    TruncatedNormal,
+    Uniform,
+)
 from pymc.distributions.discrete import Bernoulli
 from pymc.distributions.logprob import (
     _get_scaling,
@@ -235,6 +241,21 @@ def test_logp_helper():
     np.testing.assert_almost_equal(x_logp.eval(), sp.norm(0, 1).logpdf([0, 1]))
 
 
+def test_logp_helper_derived_rv():
+    assert np.isclose(
+        logp(at.exp(Normal.dist()), 5).eval(),
+        logp(LogNormal.dist(), 5).eval(),
+    )
+
+
+def test_logp_helper_exceptions():
+    with pytest.raises(TypeError, match="When RV is not a pure distribution"):
+        logp(at.exp(Normal.dist()), [1, 2])
+
+    with pytest.raises(NotImplementedError, match="PyMC could not infer logp of input variable"):
+        logp(at.cos(Normal.dist()), 1)
+
+
 def test_logcdf_helper():
     value = at.vector("value")
     x = Normal.dist(0, 1)