Merge remote-tracking branch 'upstream/main' into feature/gp-cov-type-hints

Author: Joseph Hall

.github/workflows/tests.yml

@@ -413,7 +413,7 @@ jobs:
         floatx: [float32]
         python-version: ["3.11"]
         test-subset:
-        - tests/sampling/test_mcmc.py tests/ode/test_ode.py tests/ode/test_utils.py
+        - tests/sampling/test_mcmc.py tests/ode/test_ode.py tests/ode/test_utils.py tests/distributions/test_transform.py
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:

pymc/distributions/continuous.py

@@ -57,6 +57,7 @@
 from pytensor.tensor.var import TensorConstant
 
 from pymc.logprob.abstract import _logcdf_helper, _logprob_helper
+from pymc.logprob.basic import icdf
 
 try:
     from polyagamma import polyagamma_cdf, polyagamma_pdf, random_polyagamma
@@ -856,6 +857,11 @@ def logcdf(value, loc, sigma):
             msg="sigma > 0",
         )
 
+    def icdf(value, loc, sigma):
+        res = icdf(Normal.dist(loc, sigma), (value + 1.0) / 2.0)
+        res = check_icdf_value(res, value)
+        return res
+
 
 class WaldRV(RandomVariable):
     name = "wald"
@@ -1714,12 +1720,17 @@ def logcdf(value, mu, sigma):
             -np.inf,
             normal_lcdf(mu, sigma, pt.log(value)),
         )
+
         return check_parameters(
             res,
             sigma > 0,
             msg="sigma > 0",
         )
 
+    def icdf(value, mu, sigma):
+        res = pt.exp(icdf(Normal.dist(mu, sigma), value))
+        return res
+
 
 Lognormal = LogNormal
 
@@ -2121,6 +2132,15 @@ def logcdf(value, loc, beta):
             msg="beta > 0",
         )
 
+    def icdf(value, loc, beta):
+        res = loc + beta * pt.tan(np.pi * (value) / 2.0)
+        res = check_icdf_value(res, value)
+        return check_parameters(
+            res,
+            beta > 0,
+            msg="beta > 0",
+        )
+
 
 class Gamma(PositiveContinuous):
     r"""
@@ -2526,6 +2546,16 @@ def logp(value, alpha, beta):
             msg="alpha > 0, beta > 0",
         )
 
+    def icdf(value, alpha, beta):
+        res = beta * (-pt.log(1 - value)) ** (1 / alpha)
+        res = check_icdf_value(res, value)
+        return check_parameters(
+            res,
+            alpha > 0,
+            beta > 0,
+            msg="alpha > 0, beta > 0",
+        )
+
 
 class HalfStudentTRV(RandomVariable):
     name = "halfstudentt"
@@ -3069,6 +3099,20 @@ def logcdf(value, lower, c, upper):
             msg="lower <= c <= upper",
         )
 
+    def icdf(value, lower, c, upper):
+        res = pt.switch(
+            pt.lt(value, ((c - lower) / (upper - lower))),
+            lower + np.sqrt((upper - lower) * (c - lower) * value),
+            upper - np.sqrt((upper - lower) * (upper - c) * (1 - value)),
+        )
+        res = check_icdf_value(res, value)
+        return check_parameters(
+            res,
+            lower <= c,
+            c <= upper,
+            msg="lower <= c <= upper",
+        )
+
 
 @_default_transform.register(Triangular)
 def triangular_default_transform(op, rv):
@@ -3157,6 +3201,15 @@ def logcdf(value, mu, beta):
             msg="beta > 0",
         )
 
+    def icdf(value, mu, beta):
+        res = mu - beta * pt.log(-pt.log(value))
+        res = check_icdf_value(res, value)
+        return check_parameters(
+            res,
+            beta > 0,
+            msg="beta > 0",
+        )
+
 
 class RiceRV(RandomVariable):
     name = "rice"
@@ -3713,6 +3766,15 @@ def logcdf(value, mu, sigma):
             msg="sigma > 0",
         )
 
+    def icdf(value, mu, sigma):
+        res = sigma * -pt.log(2.0 * pt.erfcinv(value) ** 2) + mu
+        res = check_icdf_value(res, value)
+        return check_parameters(
+            res,
+            sigma > 0,
+            msg="sigma > 0",
+        )
+
 
 class PolyaGammaRV(RandomVariable):
     """Polya-Gamma random variable."""

pymc/distributions/multivariate.py

@@ -132,6 +132,9 @@ def quaddist_matrix(cov=None, chol=None, tau=None, lower=True, *args, **kwargs):
         chol = pt.as_tensor_variable(chol)
         if chol.ndim != 2:
             raise ValueError("chol must be two dimensional.")
+
+        # tag as lower triangular to enable pytensor rewrites of chol(l.l') -> l
+        chol.tag.lower_triangular = True
         cov = chol.dot(chol.T)
 
     return cov

pymc/logprob/transforms.py

@@ -448,7 +448,8 @@ def measurable_transform_logcdf(op: MeasurableTransform, value, *inputs, **kwarg
 
     backward_value = op.transform_elemwise.backward(value, *other_inputs)
 
-    # Some transformations, like squaring may produce multiple backward values
+    # Fail if transformation is not injective
+    # A TensorVariable is returned in 1-to-1 inversions, and a tuple in 1-to-many
     if isinstance(backward_value, tuple):
         raise NotImplementedError
 
@@ -469,6 +470,11 @@ def measurable_transform_icdf(op: MeasurableTransform, value, *inputs, **kwargs)
     input_icdf = _icdf_helper(measurable_input, value)
     icdf = op.transform_elemwise.forward(input_icdf, *other_inputs)
 
+    # Fail if transformation is not injective
+    # A TensorVariable is returned in 1-to-1 inversions, and a tuple in 1-to-many
+    if isinstance(op.transform_elemwise.backward(icdf, *other_inputs), tuple):
+        raise NotImplementedError
+
     return icdf
 
 
@@ -958,8 +964,10 @@ class SimplexTransform(RVTransform):
     name = "simplex"
 
     def forward(self, value, *inputs):
+        value = pt.as_tensor(value)
         log_value = pt.log(value)
-        shift = pt.sum(log_value, -1, keepdims=True) / value.shape[-1]
+        N = value.shape[-1].astype(value.dtype)
+        shift = pt.sum(log_value, -1, keepdims=True) / N
         return log_value[..., :-1] - shift
 
     def backward(self, value, *inputs):
@@ -968,7 +976,9 @@ def backward(self, value, *inputs):
         return exp_value_max / pt.sum(exp_value_max, -1, keepdims=True)
 
     def log_jac_det(self, value, *inputs):
+        value = pt.as_tensor(value)
         N = value.shape[-1] + 1
+        N = N.astype(value.dtype)
         sum_value = pt.sum(value, -1, keepdims=True)
         value_sum_expanded = value + sum_value
         value_sum_expanded = pt.concatenate([value_sum_expanded, pt.zeros(sum_value.shape)], -1)

pymc/math.py

@@ -443,11 +443,17 @@ def expand_packed_triangular(n, packed, lower=True, diagonal_only=False):
     elif lower:
         out = pt.zeros((n, n), dtype=pytensor.config.floatX)
         idxs = np.tril_indices(n)
-        return pt.set_subtensor(out[idxs], packed)
+        # tag as lower triangular to enable pytensor rewrites
+        out = pt.set_subtensor(out[idxs], packed)
+        out.tag.lower_triangular = True
+        return out
     elif not lower:
         out = pt.zeros((n, n), dtype=pytensor.config.floatX)
         idxs = np.triu_indices(n)
-        return pt.set_subtensor(out[idxs], packed)
+        # tag as upper triangular to enable pytensor rewrites
+        out = pt.set_subtensor(out[idxs], packed)
+        out.tag.upper_triangular = True
+        return out
 
 
 class BatchedDiag(Op):

pymc/model.py

@@ -75,6 +75,7 @@
     hessian,
     inputvars,
     replace_rvs_by_values,
+    rewrite_pregrad,
 )
 from pymc.util import (
     UNSET,
@@ -381,6 +382,8 @@ def __init__(
             self._extra_vars_shared[var.name] = shared
             givens.append((var, shared))
 
+        cost = rewrite_pregrad(cost)
+
         if compute_grads:
             grads = pytensor.grad(cost, grad_vars, disconnected_inputs="ignore")
             for grad_wrt, var in zip(grads, grad_vars):
@@ -824,6 +827,7 @@ def dlogp(
                     )
 
         cost = self.logp(jacobian=jacobian)
+        cost = rewrite_pregrad(cost)
         return gradient(cost, value_vars)
 
     def d2logp(
@@ -862,6 +866,7 @@ def d2logp(
                     )
 
         cost = self.logp(jacobian=jacobian)
+        cost = rewrite_pregrad(cost)
         return hessian(cost, value_vars)
 
     @property

pymc/pytensorf.py

@@ -1228,3 +1228,10 @@ def constant_fold(
     return tuple(
         folded_x.data if isinstance(folded_x, Constant) else folded_x for folded_x in folded_xs
     )
+
+
+def rewrite_pregrad(graph):
+    """Apply simplifying or stabilizing rewrites to graph that are safe to use
+    pre-grad.
+    """
+    return rewrite_graph(graph, include=("canonicalize", "stabilize"))

tests/distributions/test_continuous.py

@@ -207,6 +207,12 @@ def test_triangular(self):
             lambda value, c, lower, upper: st.triang.logcdf(value, c - lower, lower, upper - lower),
             skip_paramdomain_outside_edge_test=True,
         )
+        check_icdf(
+            pm.Triangular,
+            {"lower": -Rplusunif, "c": Runif, "upper": Rplusunif},
+            lambda q, c, lower, upper: st.triang.ppf(q, c - lower, lower, upper - lower),
+            skip_paramdomain_outside_edge_test=True,
+        )
 
         # Custom logp/logcdf check for values outside of domain
         valid_dist = pm.Triangular.dist(lower=0, upper=1, c=0.9, size=2)
@@ -299,6 +305,11 @@ def test_half_normal(self):
             {"sigma": Rplus},
             lambda value, sigma: st.halfnorm.logcdf(value, scale=sigma),
         )
+        check_icdf(
+            pm.HalfNormal,
+            {"sigma": Rplus},
+            lambda q, sigma: st.halfnorm.ppf(q, scale=sigma),
+        )
 
     def test_chisquared_logp(self):
         check_logp(
@@ -502,6 +513,21 @@ def test_lognormal(self):
             {"mu": R, "sigma": Rplusbig},
             lambda value, mu, sigma: st.lognorm.logcdf(value, sigma, 0, np.exp(mu)),
         )
+        check_icdf(
+            pm.LogNormal,
+            {"mu": R, "tau": Rplusbig},
+            lambda q, mu, tau: floatX(st.lognorm.ppf(q, tau**-0.5, 0, np.exp(mu))),
+        )
+        # Because we exponentiate the normal quantile function, setting sigma >= 9.5
+        # return extreme values that results in relative errors above 4 digits
+        # we circumvent it by keeping it below or equal to 9.
+        custom_rplusbig = Domain([0, 0.5, 0.9, 0.99, 1, 1.5, 2, 9, np.inf])
+        check_icdf(
+            pm.LogNormal,
+            {"mu": R, "sigma": custom_rplusbig},
+            lambda q, mu, sigma: floatX(st.lognorm.ppf(q, sigma, 0, np.exp(mu))),
+            decimal=select_by_precision(float64=4, float32=3),
+        )
 
     def test_studentt_logp(self):
         check_logp(
@@ -567,6 +593,9 @@ def test_half_cauchy(self):
             {"beta": Rplusbig},
             lambda value, beta: st.halfcauchy.logcdf(value, scale=beta),
         )
+        check_icdf(
+            pm.HalfCauchy, {"beta": Rplusbig}, lambda q, beta: st.halfcauchy.ppf(q, scale=beta)
+        )
 
     def test_gamma_logp(self):
         check_logp(
@@ -681,6 +710,13 @@ def test_weibull_logcdf(self):
             lambda value, alpha, beta: st.exponweib.logcdf(value, 1, alpha, scale=beta),
         )
 
+    def test_weibull_icdf(self):
+        check_icdf(
+            pm.Weibull,
+            {"alpha": Rplusbig, "beta": Rplusbig},
+            lambda q, alpha, beta: st.exponweib.ppf(q, 1, alpha, scale=beta),
+        )
+
     def test_half_studentt(self):
         # this is only testing for nu=1 (halfcauchy)
         check_logp(
@@ -757,6 +793,11 @@ def test_gumbel(self):
             {"mu": R, "beta": Rplusbig},
             lambda value, mu, beta: st.gumbel_r.logcdf(value, loc=mu, scale=beta),
         )
+        check_icdf(
+            pm.Gumbel,
+            {"mu": R, "beta": Rplusbig},
+            lambda q, mu, beta: st.gumbel_r.ppf(q, loc=mu, scale=beta),
+        )
 
     def test_logistic(self):
         check_logp(
@@ -840,6 +881,13 @@ def test_moyal_logcdf(self):
         if pytensor.config.floatX == "float32":
             raise Exception("Flaky test: It passed this time, but XPASS is not allowed.")
 
+    def test_moyal_icdf(self):
+        check_icdf(
+            pm.Moyal,
+            {"mu": R, "sigma": Rplusbig},
+            lambda q, mu, sigma: floatX(st.moyal.ppf(q, mu, sigma)),
+        )
+
     def test_interpolated(self):
         for mu in R.vals:
             for sigma in Rplus.vals: