Added aeppl based log-likelihood graph generation and aeppl based transforms

Co-authored-by:
Ricardo Vieira <[email protected]>

Author: kc611

conda-envs/environment-dev-py37.yml

@@ -4,6 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
+- aeppl>=0.0.13
 - aesara>=2.2.2
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-dev-py38.yml

@@ -4,6 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
+- aeppl>=0.0.13
 - aesara>=2.2.2
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-dev-py39.yml

@@ -4,6 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
+- aeppl>=0.0.13
 - aesara>=2.2.2
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-test-py37.yml

@@ -4,6 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
+- aeppl>=0.0.13
 - aesara>=2.2.2
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-test-py38.yml

@@ -4,6 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
+- aeppl>=0.0.13
 - aesara>=2.2.2
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/environment-test-py39.yml

@@ -4,6 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
+- aeppl>=0.0.13
 - aesara>=2.2.2
 - arviz>=0.11.4
 - cachetools

conda-envs/windows-environment-dev-py38.yml

@@ -4,6 +4,7 @@ channels:
 - defaults
 dependencies:
  # base dependencies (see install guide for Windows)
+- aeppl>=0.0.13
 - aesara>=2.2.2
 - arviz>=0.11.4
 - cachetools>=4.2.1

conda-envs/windows-environment-test-py38.yml

@@ -4,6 +4,7 @@ channels:
 - defaults
 dependencies:
  # base dependencies (see install guide for Windows)
+- aeppl>=0.0.13
 - aesara>=2.2.2
 - arviz>=0.11.2
 - cachetools

pymc/aesaraf.py

@@ -377,7 +377,7 @@ def transform_replacements(var, replacements):
             # potential replacements
             return [rv_value_var]
 
-        trans_rv_value = transform.backward(rv_var, rv_value_var)
+        trans_rv_value = transform.backward(rv_value_var, *rv_var.owner.inputs)
         replacements[var] = trans_rv_value
 
         # Walk the transformed variable and make replacements

pymc/bart/bart.py

@@ -12,8 +12,10 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import aesara.tensor as at
 import numpy as np
 
+from aeppl.logprob import _logprob
 from aesara.tensor.random.op import RandomVariable, default_shape_from_params
 from pandas import DataFrame, Series
 
@@ -146,6 +148,20 @@ def __new__(
     def dist(cls, *params, **kwargs):
         return super().dist(params, **kwargs)
 
+    def logp(x, *inputs):
+        """Calculate log probability.
+
+        Parameters
+        ----------
+        x: numeric, TensorVariable
+            Value for which log-probability is calculated.
+
+        Returns
+        -------
+        TensorVariable
+        """
+        return at.zeros_like(x)
+
 
 def preprocess_XY(X, Y):
     if isinstance(Y, (Series, DataFrame)):
@@ -156,3 +172,10 @@ def preprocess_XY(X, Y):
     Y = Y.astype(float)
     X = X.astype(float)
     return X, Y
+
+
+@_logprob.register(BARTRV)
+def logp(op, value_var, *dist_params, **kwargs):
+    _dist_params = dist_params[3:]
+    value_var = value_var[0]
+    return BART.logp(value_var, *_dist_params)

pymc/distributions/__init__.py

@@ -14,9 +14,9 @@
 
 from pymc.distributions.logprob import (  # isort:skip
     _logcdf,
-    _logp,
     logcdf,
     logp,
+    logcdfpt,
     logp_transform,
     logpt,
     logpt_sum,
@@ -193,7 +193,6 @@
     "PolyaGamma",
     "logpt",
     "logp",
-    "_logp",
     "logp_transform",
     "logcdf",
     "_logcdf",

pymc/distributions/bound.py

@@ -14,12 +14,12 @@
 
 import numpy as np
 
+from aeppl.logprob import logprob
 from aesara.tensor import as_tensor_variable
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.var import TensorVariable
 
 from pymc.aesaraf import floatX, intX
-from pymc.distributions import _logp
 from pymc.distributions.continuous import BoundedContinuous
 from pymc.distributions.dist_math import bound
 from pymc.distributions.distribution import Continuous, Discrete
@@ -46,7 +46,7 @@ def rng_fn(cls, rng, distribution, lower, upper, size):
 
 class _ContinuousBounded(BoundedContinuous):
     rv_op = boundrv
-    bound_args_indices = [1, 2]
+    bound_args_indices = [4, 5]
 
     def logp(value, distribution, lower, upper):
         """
@@ -67,7 +67,7 @@ def logp(value, distribution, lower, upper):
         -------
         TensorVariable
         """
-        logp = _logp(distribution.owner.op, value, {}, *distribution.owner.inputs[3:])
+        logp = logprob(distribution, value)
         return bound(logp, (value >= lower), (value <= upper))
 
 
@@ -107,7 +107,7 @@ def logp(value, distribution, lower, upper):
         -------
         TensorVariable
         """
-        logp = _logp(distribution.owner.op, value, {}, *distribution.owner.inputs[3:])
+        logp = logprob(distribution, value)
         return bound(logp, (value >= lower), (value <= upper))
 
 
@@ -166,6 +166,7 @@ def __new__(
                 raise ValueError("Given dims do not exist in model coordinates.")
 
         lower, upper, initval = cls._set_values(lower, upper, size, shape, initval)
+        distribution.tag.ignore_logprob = True
 
         if isinstance(distribution.owner.op, Continuous):
             res = _ContinuousBounded(
@@ -200,7 +201,7 @@ def dist(
 
         cls._argument_checks(distribution, **kwargs)
         lower, upper, initval = cls._set_values(lower, upper, size, shape, initval=None)
-
+        distribution.tag.ignore_logprob = True
         if isinstance(distribution.owner.op, Continuous):
             res = _ContinuousBounded.dist(
                 [distribution, lower, upper],