Update PyMC dependency

Author: ricardoV94

.github/workflows/test.yml

@@ -26,7 +26,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      AESARA_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=native'
+      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=native'
     defaults:
       run:
         shell: bash -l {0}
@@ -91,7 +91,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      AESARA_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=core2'
+      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=core2'
     defaults:
       run:
         shell: cmd

pymc_experimental/distributions/continuous.py

@@ -21,14 +21,14 @@
 
 from typing import List, Tuple, Union
 
-import aesara.tensor as at
 import numpy as np
-from aesara.tensor.random.op import RandomVariable
-from aesara.tensor.var import TensorVariable
-from pymc.aesaraf import floatX
+import pytensor.tensor as pt
 from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import Continuous
 from pymc.distributions.shape_utils import rv_size_is_none
+from pymc.pytensorf import floatX
+from pytensor.tensor.random.op import RandomVariable
+from pytensor.tensor.var import TensorVariable
 from scipy import stats
 
 
@@ -144,9 +144,9 @@ def dist(cls, mu=0, sigma=1, xi=0, scipy=False, **kwargs):
         # If SciPy, use its parametrization, otherwise convert to standard
         if scipy:
             xi = -xi
-        mu = at.as_tensor_variable(floatX(mu))
-        sigma = at.as_tensor_variable(floatX(sigma))
-        xi = at.as_tensor_variable(floatX(xi))
+        mu = pt.as_tensor_variable(floatX(mu))
+        sigma = pt.as_tensor_variable(floatX(sigma))
+        xi = pt.as_tensor_variable(floatX(xi))
 
         return super().dist([mu, sigma, xi], **kwargs)
 
@@ -159,26 +159,26 @@ def logp(value, mu, sigma, xi):
         ----------
         value: numeric
             Value(s) for which log-probability is calculated. If the log probabilities for multiple
-            values are desired the values must be provided in a numpy array or Aesara tensor
+            values are desired the values must be provided in a numpy array or Pytensor tensor
 
         Returns
         -------
         TensorVariable
         """
         scaled = (value - mu) / sigma
 
-        logp_expression = at.switch(
-            at.isclose(xi, 0),
-            -at.log(sigma) - scaled - at.exp(-scaled),
-            -at.log(sigma)
-            - ((xi + 1) / xi) * at.log1p(xi * scaled)
-            - at.pow(1 + xi * scaled, -1 / xi),
+        logp_expression = pt.switch(
+            pt.isclose(xi, 0),
+            -pt.log(sigma) - scaled - pt.exp(-scaled),
+            -pt.log(sigma)
+            - ((xi + 1) / xi) * pt.log1p(xi * scaled)
+            - pt.pow(1 + xi * scaled, -1 / xi),
         )
 
-        logp = at.switch(at.gt(1 + xi * scaled, 0.0), logp_expression, -np.inf)
+        logp = pt.switch(pt.gt(1 + xi * scaled, 0.0), logp_expression, -np.inf)
 
         return check_parameters(
-            logp, sigma > 0, at.and_(xi > -1, xi < 1), msg="sigma > 0 or -1 < xi < 1"
+            logp, sigma > 0, pt.and_(xi > -1, xi < 1), msg="sigma > 0 or -1 < xi < 1"
         )
 
     def logcdf(value, mu, sigma, xi):
@@ -198,21 +198,21 @@ def logcdf(value, mu, sigma, xi):
         TensorVariable
         """
         scaled = (value - mu) / sigma
-        logc_expression = at.switch(
-            at.isclose(xi, 0), -at.exp(-scaled), -at.pow(1 + xi * scaled, -1 / xi)
+        logc_expression = pt.switch(
+            pt.isclose(xi, 0), -pt.exp(-scaled), -pt.pow(1 + xi * scaled, -1 / xi)
         )
 
-        logc = at.switch(1 + xi * (value - mu) / sigma > 0, logc_expression, -np.inf)
+        logc = pt.switch(1 + xi * (value - mu) / sigma > 0, logc_expression, -np.inf)
 
         return check_parameters(
-            logc, sigma > 0, at.and_(xi > -1, xi < 1), msg="sigma > 0 or -1 < xi < 1"
+            logc, sigma > 0, pt.and_(xi > -1, xi < 1), msg="sigma > 0 or -1 < xi < 1"
         )
 
     def moment(rv, size, mu, sigma, xi):
         r"""
         Using the mode, as the mean can be infinite when :math:`\xi > 1`
         """
-        mode = at.switch(at.isclose(xi, 0), mu, mu + sigma * (at.pow(1 + xi, -xi) - 1) / xi)
+        mode = pt.switch(pt.isclose(xi, 0), mu, mu + sigma * (pt.pow(1 + xi, -xi) - 1) / xi)
         if not rv_size_is_none(size):
-            mode = at.full(size, mode)
+            mode = pt.full(size, mode)
         return mode

pymc_experimental/distributions/histogram_utils.py

@@ -99,15 +99,15 @@ def histogram_approximation(name, dist, *, observed, **h_kwargs):
     ----------
     name : str
         Name for the Potential
-    dist : aesara.tensor.var.TensorVariable
+    dist : pytensor.tensor.var.TensorVariable
         The output of pm.Distribution.dist()
     observed : ArrayLike
         Observed value to construct a histogram. Histogram is computed over 0th axis.
         Dask is supported.
 
     Returns
     -------
-    aesara.tensor.var.TensorVariable
+    pytensor.tensor.var.TensorVariable
         Potential
 
     Examples

pymc_experimental/gp/latent_approx.py

@@ -13,9 +13,9 @@
 #   limitations under the License.
 
 
-import aesara.tensor as at
 import numpy as np
 import pymc as pm
+import pytensor.tensor as pt
 from pymc.gp.util import JITTER_DEFAULT, cholesky, solve_lower, solve_upper, stabilize
 
 
@@ -42,7 +42,7 @@ def _build_prior(self, name, X, Xu, jitter=JITTER_DEFAULT, **kwargs):
         u = pm.Deterministic(name + "_u", L @ v)
 
         Kfu = self.cov_func(X, Xu)
-        Kuuiu = solve_upper(at.transpose(L), solve_lower(L, u))
+        Kuuiu = solve_upper(pt.transpose(L), solve_lower(L, u))
 
         return pm.Deterministic(name, mu + Kfu @ Kuuiu), Kuuiu, L
 
@@ -62,8 +62,8 @@ def prior(self, name, X, Xu=None, jitter=JITTER_DEFAULT, **kwargs):
     def _build_conditional(self, name, Xnew, Xu, L, Kuuiu, jitter, **kwargs):
         Ksu = self.cov_func(Xnew, Xu)
         mu = self.mean_func(Xnew) + Ksu @ Kuuiu
-        tmp = solve_lower(L, at.transpose(Ksu))
-        Qss = at.transpose(tmp) @ tmp  # Qss = tt.dot(tt.dot(Ksu, tt.nlinalg.pinv(Kuu)), Ksu.T)
+        tmp = solve_lower(L, pt.transpose(Ksu))
+        Qss = pt.transpose(tmp) @ tmp  # Qss = tt.dot(tt.dot(Ksu, tt.nlinalg.pinv(Kuu)), Ksu.T)
         Kss = self.cov_func(Xnew)
         Lss = cholesky(stabilize(Kss - Qss, jitter))
         return mu, Lss
@@ -100,42 +100,42 @@ def prior(self, name, X, **kwargs):
         return f
 
     def _generate_basis(self, X, L):
-        indices = at.arange(1, self.M + 1)
-        m1 = (np.pi / (2.0 * L)) * at.tile(L + X, self.M)
-        m2 = at.diag(indices)
-        Phi = at.sin(m1 @ m2) / at.sqrt(L)
+        indices = pt.arange(1, self.M + 1)
+        m1 = (np.pi / (2.0 * L)) * pt.tile(L + X, self.M)
+        m2 = pt.diag(indices)
+        Phi = pt.sin(m1 @ m2) / pt.sqrt(L)
         omega = (np.pi * indices) / (2.0 * L)
         return Phi, omega
 
     def _build_prior(self, name, X, **kwargs):
         n_obs = np.shape(X)[0]
 
         # standardize input scale
-        X = at.as_tensor_variable(X)
-        Xmu = at.mean(X, axis=0)
-        Xsd = at.std(X, axis=0)
+        X = pt.as_tensor_variable(X)
+        Xmu = pt.mean(X, axis=0)
+        Xsd = pt.std(X, axis=0)
         Xz = (X - Xmu) / Xsd
 
         # define L using Xz and c
-        La = at.abs(at.min(Xz))  # .eval()?
-        Lb = at.max(Xz)
-        L = self.c * at.max([La, Lb])
+        La = pt.abs(pt.min(Xz))  # .eval()?
+        Lb = pt.max(Xz)
+        L = self.c * pt.max([La, Lb])
 
         # make basis and omega, spectral density
         Phi, omega = self._generate_basis(Xz, L)
         scale, ls, spectral_density = self._validate_cov_func(self.cov_func)
         spd = scale * spectral_density(omega, ls / Xsd).flatten()
 
         beta = pm.Normal(f"{name}_coeffs_", size=self.M)
-        f = pm.Deterministic(name, self.mean_func(X) + at.dot(Phi * at.sqrt(spd), beta))
+        f = pm.Deterministic(name, self.mean_func(X) + pt.dot(Phi * pt.sqrt(spd), beta))
         return f, Phi, L, spd, beta, Xmu, Xsd
 
     def _build_conditional(self, Xnew, Xmu, Xsd, L, beta):
         Xnewz = (Xnew - Xmu) / Xsd
         Phi, omega = self._generate_basis(Xnewz, L)
         scale, ls, spectral_density = self._validate_cov_func(self.cov_func)
         spd = scale * spectral_density(omega, ls / Xsd).flatten()
-        return self.mean_func(Xnew) + at.dot(Phi * at.sqrt(spd), beta)
+        return self.mean_func(Xnew) + pt.dot(Phi * pt.sqrt(spd), beta)
 
     def conditional(self, name, Xnew):
         # warn about extrapolation
@@ -147,15 +147,15 @@ class ExpQuad(pm.gp.cov.ExpQuad):
     @staticmethod
     def spectral_density(omega, ls):
         # univariate spectral denisty, implement multi
-        return at.sqrt(2 * np.pi) * ls * at.exp(-0.5 * ls**2 * omega**2)
+        return pt.sqrt(2 * np.pi) * ls * pt.exp(-0.5 * ls**2 * omega**2)
 
 
 class Matern52(pm.gp.cov.Matern52):
     @staticmethod
     def spectral_density(omega, ls):
         # univariate spectral denisty, implement multi
         # https://arxiv.org/pdf/1611.06740.pdf
-        lam = at.sqrt(5) * (1.0 / ls)
+        lam = pt.sqrt(5) * (1.0 / ls)
         return (16.0 / 3.0) * lam**5 * (1.0 / (lam**2 + omega**2) ** 3)
 
 
@@ -165,7 +165,7 @@ def spectral_density(omega, ls):
         # univariate spectral denisty, implement multi
         # https://arxiv.org/pdf/1611.06740.pdf
         lam = np.sqrt(3.0) * (1.0 / ls)
-        return 4.0 * lam**3 * (1.0 / at.square(lam**2 + omega**2))
+        return 4.0 * lam**3 * (1.0 / pt.square(lam**2 + omega**2))
 
 
 class Matern12(pm.gp.cov.Matern12):
@@ -193,7 +193,7 @@ def __init__(
     def _build_prior(self, name, X, jitter=1e-6, **kwargs):
         mu = self.mean_func(X)
         Kxx = pm.gp.util.stabilize(self.cov_func(X), jitter)
-        vals, vecs = at.linalg.eigh(Kxx)
+        vals, vecs = pt.linalg.eigh(Kxx)
         ## NOTE: REMOVED PRECISION CUTOFF
         if self.variance_limit is None:
             n_eigs = self.n_eigs
@@ -204,7 +204,7 @@ def _build_prior(self, name, X, jitter=1e-6, **kwargs):
                 n_eigs = ((vals[::-1].cumsum() / vals.sum()) > self.variance_limit).nonzero()[0][0]
         U = vecs[:, -n_eigs:]
         s = vals[-n_eigs:]
-        basis = U * at.sqrt(s)
+        basis = U * pt.sqrt(s)
 
         coefs_raw = pm.Normal(f"_gp_{name}_coefs", mu=0, sigma=1, size=n_eigs)
         # weight = pm.HalfNormal(f"_gp_{name}_sd")
@@ -222,7 +222,7 @@ def prior(self, name, X, jitter=1e-6, **kwargs):
     def _build_conditional(self, Xnew, X, f, U, s, jitter):
         Kxs = self.cov_func(X, Xnew)
         Kss = self.cov_func(Xnew)
-        Kxxpinv = U @ at.diag(1.0 / s) @ U.T
+        Kxxpinv = U @ pt.diag(1.0 / s) @ U.T
         mus = Kxs.T @ Kxxpinv @ f
         K = Kss - Kxs.T @ Kxxpinv @ Kxs
         L = pm.gp.util.cholesky(pm.gp.util.stabilize(K, jitter))

pymc_experimental/inference/pathfinder.py

@@ -105,7 +105,7 @@ def fit_pathfinder(
     init_position_dict = model.initial_point()
     init_position = [init_position_dict[rv] for rv in rvs]
 
-    new_logprob, new_input = pm.aesaraf.join_nonshared_inputs(
+    new_logprob, new_input = pm.pytensorf.join_nonshared_inputs(
         init_position_dict, (model.logp(),), model.value_vars, ()
     )
 

pymc_experimental/marginal_model.py

@@ -1,22 +1,21 @@
 import warnings
 from typing import Sequence, Tuple, Union
 
-import aesara.tensor as at
 import numpy as np
-from aeppl import factorized_joint_logprob
-from aeppl.abstract import _get_measurable_outputs
-from aeppl.logprob import _logprob
-from aesara import Mode
-from aesara.compile import SharedVariable
-from aesara.compile.builders import OpFromGraph
-from aesara.graph import Constant, FunctionGraph, ancestors, clone_replace
-from aesara.scan import map as scan_map
-from aesara.tensor import TensorVariable
-from aesara.tensor.elemwise import Elemwise
+import pytensor.tensor as pt
 from pymc import SymbolicRandomVariable
-from pymc.aesaraf import constant_fold, inputvars
 from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
+from pymc.logprob.abstract import _get_measurable_outputs, _logprob
+from pymc.logprob.joint_logprob import factorized_joint_logprob
 from pymc.model import Model
+from pymc.pytensorf import constant_fold, inputvars
+from pytensor import Mode
+from pytensor.compile import SharedVariable
+from pytensor.compile.builders import OpFromGraph
+from pytensor.graph import Constant, FunctionGraph, ancestors, clone_replace
+from pytensor.scan import map as scan_map
+from pytensor.tensor import TensorVariable
+from pytensor.tensor.elemwise import Elemwise
 
 __all__ = ["MarginalModel"]
 
@@ -381,7 +380,7 @@ def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> Tuple[int, ...]:
         return (0, 1)
     elif isinstance(op, Categorical):
         p_param = rv.owner.inputs[3]
-        return tuple(range(at.get_vector_length(p_param)))
+        return tuple(range(pt.get_vector_length(p_param)))
     elif isinstance(op, DiscreteUniform):
         lower, upper = constant_fold(rv.owner.inputs[3:])
         return tuple(range(lower, upper + 1))
@@ -437,8 +436,8 @@ def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
     # PyMC does not allow RVs in the logp graph, even if we are just using the shape
     marginalized_rv_shape = constant_fold(tuple(marginalized_rv.shape))
     marginalized_rv_domain = get_domain_of_finite_discrete_rv(marginalized_rv)
-    marginalized_rv_domain_tensor = at.swapaxes(
-        at.full(
+    marginalized_rv_domain_tensor = pt.swapaxes(
+        pt.full(
             (*marginalized_rv_shape, len(marginalized_rv_domain)),
             marginalized_rv_domain,
             dtype=marginalized_rv.dtype,
@@ -459,7 +458,7 @@ def finite_discrete_marginal_rv_logp(op, values, *inputs, **kwargs):
         ]
     else:
         # Make sure this is rewrite is registered
-        from pymc.aesaraf import local_remove_check_parameter
+        from pymc.pytensorf import local_remove_check_parameter
 
         def logp_fn(marginalized_rv_const, *non_sequences):
             return joint_logp_op(marginalized_rv_const, *non_sequences)
@@ -471,7 +470,7 @@ def logp_fn(marginalized_rv_const, *non_sequences):
             mode=Mode().including("local_remove_check_parameter"),
         )
 
-    joint_logps = at.logsumexp(joint_logps, axis=0)
+    joint_logps = pt.logsumexp(joint_logps, axis=0)
 
-    # We have to add dummy logps for the remaining value variables, otherwise AePPL will raise
-    return joint_logps, *(at.constant(0),) * (len(values) - 1)
+    # We have to add dummy logps for the remaining value variables, otherwise PyMC will raise
+    return joint_logps, *(pt.constant(0),) * (len(values) - 1)

pymc_experimental/tests/distributions/test_continuous.py

@@ -12,10 +12,11 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-# general imports
-import aesara
 import numpy as np
 import pymc as pm
+
+# general imports
+import pytensor
 import pytest
 import scipy.stats.distributions as sp
 
@@ -45,7 +46,7 @@ class TestGenExtremeClass:
     """
 
     @pytest.mark.xfail(
-        condition=(aesara.config.floatX == "float32"),
+        condition=(pytensor.config.floatX == "float32"),
         reason="PyMC underflows earlier than scipy on float32",
     )
     def test_logp(self):
@@ -62,7 +63,7 @@ def test_logp(self):
             else -np.inf,
         )
 
-        if aesara.config.floatX == "float32":
+        if pytensor.config.floatX == "float32":
             raise Exception("Flaky test: It passed this time, but XPASS is not allowed.")
 
     def test_logcdf(self):