Update dependency to aesara 2.0.8, and necessary fixes

pymc3/distributions/dist_math.py

@@ -198,7 +198,7 @@ def rho2sigma(rho):
     """
     `rho -> sigma` Aesara converter
     :math:`mu + sigma*e = mu + log(1+exp(rho))*e`"""
-    return at.nnet.softplus(rho)
+    return at.softplus(rho)
 
 
 rho2sd = rho2sigma

pymc3/distributions/transforms.py

@@ -149,7 +149,7 @@ class LogExpM1(ElemwiseTransform):
     name = "log_exp_m1"
 
     def backward(self, rv_var, rv_value):
-        return at.nnet.softplus(rv_value)
+        return at.softplus(rv_value)
 
     def forward(self, rv_var, rv_value):
         """Inverse operation of softplus.
@@ -160,7 +160,7 @@ def forward(self, rv_var, rv_value):
         return at.log(1.0 - at.exp(-rv_value)) + rv_value
 
     def jacobian_det(self, rv_var, rv_value):
-        return -at.nnet.softplus(-rv_value)
+        return -at.softplus(-rv_value)
 
 
 log_exp_m1 = LogExpM1()
@@ -191,7 +191,7 @@ def backward(self, rv_var, rv_value):
         a, b = self.param_extract_fn(rv_var)
 
         if a is not None and b is not None:
-            sigmoid_x = at.nnet.sigmoid(rv_value)
+            sigmoid_x = at.sigmoid(rv_value)
             return sigmoid_x * b + (1 - sigmoid_x) * a
         elif a is not None:
             return at.exp(rv_value) + a
@@ -215,7 +215,7 @@ def jacobian_det(self, rv_var, rv_value):
         a, b = self.param_extract_fn(rv_var)
 
         if a is not None and b is not None:
-            s = at.nnet.softplus(-rv_value)
+            s = at.softplus(-rv_value)
             return at.log(b - a) - 2 * s - rv_value
         else:
             return rv_value

pymc3/math.py

@@ -58,6 +58,7 @@
     or_,
     prod,
     sgn,
+    sigmoid,
     sin,
     sinh,
     sqr,
@@ -78,7 +79,6 @@
 
 
 from aesara.tensor.nlinalg import det, matrix_dot, matrix_inverse, trace
-from aesara.tensor.nnet import sigmoid
 from scipy.linalg import block_diag as scipy_block_diag
 
 from pymc3.aesaraf import floatX, ix_, largest_common_dtype
@@ -229,7 +229,7 @@ def log1pexp(x):
 
     This function is numerically more stable than the naive approach.
     """
-    return at.nnet.softplus(x)
+    return at.softplus(x)
 
 
 def log1mexp(x):

pymc3/variational/flows.py

@@ -390,15 +390,15 @@ def make_uw(self, u, w):
             # u_: d
             # w_: d
             wu = u.dot(w)  # .
-            mwu = -1.0 + at.nnet.softplus(wu)  # .
+            mwu = -1.0 + at.softplus(wu)  # .
             # d + (. - .) * d / .
             u_h = u + (mwu - wu) * w / ((w ** 2).sum() + 1e-10)
             return u_h, w
         else:
             # u_: bxd
             # w_: bxd
             wu = (u * w).sum(-1, keepdims=True)  # bx-
-            mwu = -1.0 + at.nnet.softplus(wu)  # bx-
+            mwu = -1.0 + at.softplus(wu)  # bx-
             # bxd + (bx- - bx-) * bxd / bx- = bxd
             u_h = u + (mwu - wu) * w / ((w ** 2).sum(-1, keepdims=True) + 1e-10)
             return u_h, w
@@ -507,7 +507,7 @@ def __init__(self, **kwargs):
 
     def make_ab(self, a, b):
         a = at.exp(a)
-        b = -a + at.nnet.softplus(b)
+        b = -a + at.softplus(b)
         return a, b
 
 

requirements.txt

@@ -1,4 +1,4 @@
-aesara>=2.0.5
+aesara>=2.0.8
 arviz>=0.11.2
 cachetools>=4.2.1
 dill