Upgrade to aesara=2.8.2 and aeppl=0.0.35

Author: Armavica

.pre-commit-config.yaml

@@ -26,8 +26,8 @@ repos:
       - types-filelock
       - types-setuptools
       - arviz
-      - aesara==2.7.9
-      - aeppl==0.0.34
+      - aesara==2.8.2
+      - aeppl==0.0.35
       always_run: true
       require_serial: true
       pass_filenames: false

conda-envs/environment-dev.yml

@@ -5,8 +5,8 @@ channels:
 - defaults
 dependencies:
 # Base dependencies
-- aeppl=0.0.34
-- aesara=2.7.9
+- aeppl=0.0.35
+- aesara=2.8.2
 - arviz>=0.12.0
 - blas
 - cachetools>=4.2.1

conda-envs/environment-test.yml

@@ -5,8 +5,8 @@ channels:
 - defaults
 dependencies:
 # Base dependencies
-- aeppl=0.0.34
-- aesara=2.7.9
+- aeppl=0.0.35
+- aesara=2.8.2
 - arviz>=0.12.0
 - blas
 - cachetools>=4.2.1

conda-envs/windows-environment-dev.yml

@@ -5,8 +5,8 @@ channels:
 - defaults
 dependencies:
 # Base dependencies (see install guide for Windows)
-- aeppl=0.0.34
-- aesara=2.7.9
+- aeppl=0.0.35
+- aesara=2.8.2
 - arviz>=0.12.0
 - blas
 - cachetools>=4.2.1

conda-envs/windows-environment-test.yml

@@ -5,8 +5,8 @@ channels:
 - defaults
 dependencies:
 # Base dependencies (see install guide for Windows)
-- aeppl=0.0.34
-- aesara=2.7.9
+- aeppl=0.0.35
+- aesara=2.8.2
 - arviz>=0.12.0
 - blas
 - cachetools>=4.2.1

pymc/aesaraf.py

@@ -37,7 +37,7 @@
 from aesara import config, scalar
 from aesara.compile.mode import Mode, get_mode
 from aesara.gradient import grad
-from aesara.graph import local_optimizer
+from aesara.graph import node_rewriter
 from aesara.graph.basic import (
     Apply,
     Constant,
@@ -875,7 +875,7 @@ def largest_common_dtype(tensors):
     return np.stack([np.ones((), dtype=dtype) for dtype in dtypes]).dtype
 
 
-@local_optimizer(tracks=[CheckParameterValue])
+@node_rewriter(tracks=[CheckParameterValue])
 def local_remove_check_parameter(fgraph, node):
     """Rewrite that removes Aeppl's CheckParameterValue
 
@@ -885,7 +885,7 @@ def local_remove_check_parameter(fgraph, node):
         return [node.inputs[0]]
 
 
-@local_optimizer(tracks=[CheckParameterValue])
+@node_rewriter(tracks=[CheckParameterValue])
 def local_check_parameter_to_ninf_switch(fgraph, node):
     if isinstance(node.op, CheckParameterValue):
         logp_expr, *logp_conds = node.inputs

pymc/distributions/continuous.py

@@ -35,7 +35,6 @@
 from aesara.tensor.math import tanh
 from aesara.tensor.random.basic import (
     BetaRV,
-    WeibullRV,
     cauchy,
     chisquare,
     exponential,
@@ -1464,7 +1463,7 @@ def dist(cls, lam, *args, **kwargs):
         lam = at.as_tensor_variable(floatX(lam))
 
         # Aesara exponential op is parametrized in terms of mu (1/lam)
-        return super().dist([at.inv(lam)], **kwargs)
+        return super().dist([at.reciprocal(lam)], **kwargs)
 
     def moment(rv, size, mu):
         if not rv_size_is_none(size):
@@ -1487,7 +1486,7 @@ def logcdf(value, mu):
         -------
         TensorVariable
         """
-        lam = at.inv(mu)
+        lam = at.reciprocal(mu)
         res = at.switch(
             at.lt(value, 0),
             -np.inf,
@@ -2313,7 +2312,7 @@ def logcdf(value, alpha, inv_beta):
         -------
         TensorVariable
         """
-        beta = at.inv(inv_beta)
+        beta = at.reciprocal(inv_beta)
         res = at.switch(
             at.lt(value, 0),
             -np.inf,
@@ -2518,8 +2517,15 @@ def logcdf(value, nu):
 
 
 # TODO: Remove this once logp for multiplication is working!
-class WeibullBetaRV(WeibullRV):
+class WeibullBetaRV(RandomVariable):
+    name = "weibull"
+    ndim_supp = 0
     ndims_params = [0, 0]
+    dtype = "floatX"
+    _print_name = ("Weibull", "\\operatorname{Weibull}")
+
+    def __call__(self, alpha, beta, size=None, **kwargs):
+        return super().__call__(alpha, beta, size=size, **kwargs)
 
     @classmethod
     def rng_fn(cls, rng, alpha, beta, size) -> np.ndarray:
@@ -2615,6 +2621,16 @@ def logcdf(value, alpha, beta):
 
         return check_parameters(res, 0 < alpha, 0 < beta, msg="alpha > 0, beta > 0")
 
+    def logp(value, alpha, beta):
+        res = (
+            at.log(alpha)
+            - at.log(beta)
+            + (alpha - 1.0) * at.log(value / beta)
+            - at.pow(value / beta, alpha)
+        )
+        res = at.switch(at.ge(value, 0.0), res, -np.inf)
+        return check_parameters(res, 0 < alpha, 0 < beta, msg="alpha > 0, beta > 0")
+
 
 class HalfStudentTRV(RandomVariable):
     name = "halfstudentt"

pymc/distributions/dist_math.py

@@ -157,7 +157,7 @@ def sigma2rho(sigma):
     """
     `sigma -> rho` Aesara converter
     :math:`mu + sigma*e = mu + log(1+exp(rho))*e`"""
-    return at.log(at.exp(at.abs_(sigma)) - 1.0)
+    return at.log(at.exp(at.abs(sigma)) - 1.0)
 
 
 def rho2sigma(rho):
@@ -213,7 +213,7 @@ def log_normal(x, mean, **kwargs):
     else:
         std = tau ** (-1)
     std += f(eps)
-    return f(c) - at.log(at.abs_(std)) - (x - mean) ** 2 / (2.0 * std**2)
+    return f(c) - at.log(at.abs(std)) - (x - mean) ** 2 / (2.0 * std**2)
 
 
 def MvNormalLogp():

pymc/distributions/logprob.py

@@ -24,7 +24,7 @@
 from aeppl.abstract import assign_custom_measurable_outputs
 from aeppl.logprob import logcdf as logcdf_aeppl
 from aeppl.logprob import logprob as logp_aeppl
-from aeppl.transforms import TransformValuesOpt
+from aeppl.transforms import TransformValuesRewrite
 from aesara.graph.basic import graph_inputs, io_toposort
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.subtensor import (
@@ -231,7 +231,7 @@ def joint_logp(
                 if original_value_var is not None and hasattr(original_value_var.tag, "transform"):
                     transform_map[value_var] = original_value_var.tag.transform
 
-    transform_opt = TransformValuesOpt(transform_map)
+    transform_opt = TransformValuesRewrite(transform_map)
     temp_logp_var_dict = factorized_joint_logprob(
         tmp_rvs_to_values,
         extra_rewrites=transform_opt,

pymc/distributions/multivariate.py

@@ -539,7 +539,7 @@ def moment(rv, size, n, p):
         n = at.shape_padright(n)
         mode = at.round(n * p)
         diff = n - at.sum(mode, axis=-1, keepdims=True)
-        inc_bool_arr = at.abs_(diff) > 0
+        inc_bool_arr = at.abs(diff) > 0
         mode = at.inc_subtensor(mode[inc_bool_arr.nonzero()], diff[inc_bool_arr.nonzero()])
         if not rv_size_is_none(size):
             output_size = at.concatenate([size, [p.shape[-1]]])