From e263ba58110baff664589b57c9d539ac91130910 Mon Sep 17 00:00:00 2001
From: Noah Shin <noahshin0911@gmail.com>
Date: Sat, 9 May 2026 17:18:02 -1000
Subject: [PATCH 1/5] Add ZeroOneInflatedBeta distribution

---
 pymc/distributions/__init__.py   |   2 +
 pymc/distributions/continuous.py | 363 ++++++++++++++++++++++++++++---
 2 files changed, 331 insertions(+), 34 deletions(-)

diff --git a/pymc/distributions/__init__.py b/pymc/distributions/__init__.py
index c578267091..87494447eb 100644
--- a/pymc/distributions/__init__.py
+++ b/pymc/distributions/__init__.py
@@ -51,6 +51,7 @@
     VonMises,
     Wald,
     Weibull,
+    ZeroOneInflatedBeta
 )
 from pymc.distributions.custom import CustomDist, DensityDist
 from pymc.distributions.discrete import (
@@ -204,5 +205,6 @@
     "ZeroInflatedBinomial",
     "ZeroInflatedNegativeBinomial",
     "ZeroInflatedPoisson",
+    "ZeroOneInflatedBeta",
     "ZeroSumNormal",
 ]
diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py
index 9c4382ffb7..77c552a865 100644
--- a/pymc/distributions/continuous.py
+++ b/pymc/distributions/continuous.py
@@ -89,7 +89,11 @@ def polyagamma_cdf(*args, **kwargs):
     normal_lcdf,
     zvalue,
 )
-from pymc.distributions.distribution import DIST_PARAMETER_TYPES, Continuous, SymbolicRandomVariable
+from pymc.distributions.distribution import (
+    DIST_PARAMETER_TYPES,
+    Continuous,
+    SymbolicRandomVariable,
+)
 from pymc.distributions.shape_utils import implicit_size_from_params, rv_size_is_none
 from pymc.distributions.transforms import _default_transform
 from pymc.math import invlogit, logdiffexp
@@ -129,6 +133,7 @@ def polyagamma_cdf(*args, **kwargs):
     "VonMises",
     "Wald",
     "Weibull",
+    "ZeroOneInflatedBeta",
 ]
 
 
@@ -169,7 +174,9 @@ def circ_cont_transform(op, rv):
 @_default_transform.register(BoundedContinuous)
 def bounded_cont_transform(op, rv, bound_args_indices=None):
     if bound_args_indices is None:
-        raise ValueError(f"Must specify bound_args_indices for {op} bounded distribution")
+        raise ValueError(
+            f"Must specify bound_args_indices for {op} bounded distribution"
+        )
 
     def transform_params(*args):
         lower, upper = None, None
@@ -377,7 +384,9 @@ def logp(value):
 
     def logcdf(value):
         return pt.switch(
-            pt.eq(value, -np.inf), -np.inf, pt.switch(pt.eq(value, np.inf), 0, pt.log(0.5))
+            pt.eq(value, -np.inf),
+            -np.inf,
+            pt.switch(pt.eq(value, np.inf), 0, pt.log(0.5)),
         )
 
 
@@ -437,7 +446,9 @@ def logp(value):
         return pt.switch(pt.lt(value, 0), -np.inf, pt.zeros_like(value))
 
     def logcdf(value):
-        return pt.switch(pt.lt(value, np.inf), -np.inf, pt.switch(pt.eq(value, np.inf), 0, -np.inf))
+        return pt.switch(
+            pt.lt(value, np.inf), -np.inf, pt.switch(pt.eq(value, np.inf), 0, -np.inf)
+        )
 
 
 class Normal(Continuous):
@@ -522,7 +533,11 @@ def support_point(rv, size, mu, sigma):
         return mu
 
     def logp(value, mu, sigma):
-        res = -0.5 * pt.pow((value - mu) / sigma, 2) - pt.log(pt.sqrt(2.0 * np.pi)) - pt.log(sigma)
+        res = (
+            -0.5 * pt.pow((value - mu) / sigma, 2)
+            - pt.log(pt.sqrt(2.0 * np.pi))
+            - pt.log(sigma)
+        )
         return check_parameters(
             res,
             sigma > 0,
@@ -679,8 +694,12 @@ def dist(
         sigma = pt.as_tensor_variable(sigma)
         mu = pt.as_tensor_variable(mu)
 
-        lower = pt.as_tensor_variable(lower) if lower is not None else pt.constant(-np.inf)
-        upper = pt.as_tensor_variable(upper) if upper is not None else pt.constant(np.inf)
+        lower = (
+            pt.as_tensor_variable(lower) if lower is not None else pt.constant(-np.inf)
+        )
+        upper = (
+            pt.as_tensor_variable(upper) if upper is not None else pt.constant(np.inf)
+        )
         return super().dist([mu, sigma, lower, upper], **kwargs)
 
     def support_point(rv, size, mu, sigma, lower, upper):
@@ -712,7 +731,9 @@ def logp(value, mu, sigma, lower, upper):
         is_lower_bounded = not (
             isinstance(lower, TensorConstant) and np.all(np.isneginf(lower.value))
         )
-        is_upper_bounded = not (isinstance(upper, TensorConstant) and np.all(np.isinf(upper.value)))
+        is_upper_bounded = not (
+            isinstance(upper, TensorConstant) and np.all(np.isinf(upper.value))
+        )
 
         if is_lower_bounded and is_upper_bounded:
             norm = log_diff_normal_cdf(mu, sigma, upper, lower)
@@ -748,7 +769,9 @@ def logcdf(value, mu, sigma, lower, upper):
         is_lower_bounded = not (
             isinstance(lower, TensorConstant) and np.all(np.isneginf(lower.value))
         )
-        is_upper_bounded = not (isinstance(upper, TensorConstant) and np.all(np.isinf(upper.value)))
+        is_upper_bounded = not (
+            isinstance(upper, TensorConstant) and np.all(np.isinf(upper.value))
+        )
 
         if is_lower_bounded:
             logcdf = pt.switch(value < lower, -np.inf, logcdf)
@@ -859,7 +882,11 @@ def support_point(rv, size, loc, sigma):
         return support_point
 
     def logp(value, loc, sigma):
-        res = -0.5 * pt.pow((value - loc) / sigma, 2) + pt.log(pt.sqrt(2.0 / np.pi)) - pt.log(sigma)
+        res = (
+            -0.5 * pt.pow((value - loc) / sigma, 2)
+            + pt.log(pt.sqrt(2.0 / np.pi))
+            - pt.log(sigma)
+        )
         res = pt.switch(pt.ge(value, loc), res, -np.inf)
         return check_parameters(
             res,
@@ -1205,7 +1232,9 @@ def logp(value, alpha, beta):
             + pt.switch(pt.eq(beta, 1.0), 0.0, (beta - 1.0) * pt.log1p(-value))
             - (pt.gammaln(alpha) + pt.gammaln(beta) - pt.gammaln(alpha + beta))
         )
-        res = pt.switch(pt.bitwise_and(pt.ge(value, 0.0), pt.le(value, 1.0)), res, -np.inf)
+        res = pt.switch(
+            pt.bitwise_and(pt.ge(value, 0.0), pt.le(value, 1.0)), res, -np.inf
+        )
         return check_parameters(
             res,
             alpha > 0,
@@ -1242,6 +1271,213 @@ def icdf(value, alpha, beta):
         )
 
 
+class ZeroOneInflatedBetaRV(SymbolicRandomVariable):
+    name = "zero_one_inflated_beta"
+    extended_signature = "[rng],[size],(),(),(),()->[rng],()"
+    _print_name = ("ZeroOneInflatedBeta", "\\operatorname{ZeroOneInflatedBeta}")
+
+    @classmethod
+    def rv_op(cls, zoi, coi, alpha, beta_, *, rng=None, size=None):
+        zoi = pt.as_tensor(zoi)
+        coi = pt.as_tensor(coi)
+        alpha = pt.as_tensor(alpha)
+        beta_ = pt.as_tensor(beta_)
+        rng = normalize_rng_param(rng)
+        size = normalize_size_param(size)
+        if rv_size_is_none(size):
+            size = implicit_size_from_params(
+                zoi, coi, alpha, beta_, ndims_params=cls.ndims_params
+            )
+
+        next_rng, u = uniform(size=size, rng=rng).owner.outputs
+
+        next_rng, beta_draws = beta(alpha, beta_, size=size, rng=next_rng).owner.outputs
+
+        draws = pt.switch(
+            pt.lt(u, zoi * (1 - coi)),
+            pt.zeros_like(beta_draws),
+            pt.switch(
+                pt.lt(u, zoi),
+                pt.ones_like(beta_draws),
+                beta_draws,
+            ),
+        )
+
+        return cls(
+            inputs=[rng, size, zoi, coi, alpha, beta_],
+            outputs=[next_rng, draws],
+        )(rng, size, zoi, coi, alpha, beta_)
+
+
+class ZeroOneInflatedBeta(UnitContinuous):
+    r"""
+    Zero-One-Inflated Beta Distribution.
+
+    The pdf of this distribution is:
+
+    .. math::
+
+        f(x \mid \text{zoi}, \text{coi}, \alpha, \beta) =
+        \begin{cases}
+            \text{zoi} \cdot (1 - \text{coi})
+                & \text{if } x = 0 \\
+            \text{zoi} \cdot \text{coi}
+                & \text{if } x = 1 \\
+            (1 - \text{zoi}) \cdot \dfrac{x^{\alpha - 1} (1 - x)^{\beta - 1}}{B(\alpha, \beta)}
+                & \text{if } x \in (0, 1)
+        \end{cases}
+
+    where :math:`B` is the Beta function, :math:`\alpha = \mu \kappa`, and :math:`\beta = (1 - \mu) \kappa`.
+
+
+    Examples
+    --------
+    Model proportion data with structural zeros and ones:
+
+    .. code-block:: python
+
+        import pymc as pm
+        import numpy as np
+
+        # Data with 30% boundary values, 40% of which are ones
+        true_zoi = 0.3
+        true_coi = 0.4
+        true_mu = 0.4
+        true_kappa = 20.0
+        n = 1000
+
+        rng = np.random.default_rng(100)
+        u = rng.random(n)
+        alpha = true_mu * true_kappa
+        beta  = (1 - true_mu) * true_kappa   
+        beta_data = rng.beta(alpha, beta, n)
+        data = np.where(u < true_zoi * (1 - true_coi), 0.0,
+               np.where(u < true_zoi, 1.0,
+               beta_data))
+
+        # Fit model
+        with pm.Model() as model:
+            zoi   = pm.Beta("zoi", alpha=1, beta=1)
+            coi   = pm.Beta("coi", alpha=1, beta=1)
+            mu    = pm.Beta("mu", alpha=2, beta=2)
+            kappa = pm.Gamma("kappa", alpha=2, beta=0.1)
+            y = pm.ZeroOneInflatedBeta(
+                "y", zoi=zoi, coi=coi, mu=mu, kappa=kappa,
+                observed=data
+            )
+            trace = pm.sample(1000, tune=1000)
+
+
+    ========  ===============================================
+    Support   :math:`y \in [0, 1]`
+    Mean      :math:`(1 - \text{zoi})\mu + \text{zoi} \cdot \text{coi}`
+    Variance  :math:`(1-\text{zoi})\left[\dfrac{\mu(1-\mu)}{\kappa+1} + \text{zoi} \cdot \text{coi} \cdot (1 - \text{coi})\right] + \text{zoi}(1-\text{zoi})\mu^2`
+    ========  ===============================================
+
+    Setting :math:`coi = 0` recovers the Zero-Inflated Beta (ZIB) distribution.
+    Setting :math:`coi = 1` recovers the One-Inflated Beta (OIB) distribution.
+
+
+    Parameters
+    ----------
+    zoi : tensor_like of float
+        Probability of a structural value at 0 or 1, :math:`0 \leq zoi \leq 1`.
+    coi : tensor_like of float
+        Conditional probability given a structural value that it is a structural 1, :math:`0 \leq coi \leq 1`.
+    alpha : tensor_like of float, optional
+        ``alpha`` > 0. If not specified, then calculated using ``mu`` and ``kappa``.
+    beta : tensor_like of float, optional
+        ``beta`` > 0. If not specified, then calculated using ``mu`` and ``kappa``.
+    mu : tensor_like of float, optional
+        Mean of the Beta component, :math:`0 < \mu < 1`.
+    kappa : tensor_like of float, optional
+        Precision of the Beta component, :math:`\kappa > 0`.
+
+    References
+    ----------
+    .. [1] Ospina, R., & Ferrari, S. L. (2010). Inflated beta distributions.
+    .. [2] Ospina, R., & Ferrari, S. L. (2012). A general class of zero-or-one inflated
+           beta regression models.
+
+    Notes
+    -----
+    ZeroOneInflatedBeta is appropriate for modeling proportion data when exact zeros
+    or ones represent distinct processes (structural zeros and ones), separate from
+    the tails of the Beta component.
+    """
+
+    rv_type = ZeroOneInflatedBetaRV
+    rv_op = ZeroOneInflatedBetaRV.rv_op
+
+    @classmethod
+    def dist(cls, zoi, coi, mu=None, kappa=None, alpha=None, beta=None, **kwargs):
+        if (
+            (alpha is not None)
+            and (beta is not None)
+            and (mu is not None or kappa is not None)
+        ):
+            warnings.warn(
+                "Both parametrizations provided. Proceeding with alpha and beta.",
+                UserWarning,
+            )
+        if (alpha is not None) and (beta is not None):
+            pass
+        elif (mu is not None) and (kappa is not None):
+            alpha = mu * kappa
+            beta = (1.0 - mu) * kappa
+        else:
+            raise ValueError(
+                "Incompatible parametrization. Must specify alpha and beta, or mu and kappa."
+            )
+        zoi = pt.as_tensor_variable(zoi)
+        coi = pt.as_tensor_variable(coi)
+        alpha = pt.as_tensor_variable(alpha)
+        beta = pt.as_tensor_variable(beta)
+        return super().dist([zoi, coi, alpha, beta], **kwargs)
+
+    def support_point(rv, size, zoi, coi, alpha, beta):
+        mean = zoi * coi + (1.0 - zoi) * (alpha / (alpha + beta))
+        if not rv_size_is_none(size):
+            mean = pt.full(size, mean)
+        return mean
+
+    def logp(value, zoi, coi, alpha, beta):
+        # Beta log-probability
+        logp_beta = (
+            pt.switch(pt.eq(alpha, 1.0), 0.0, (alpha - 1.0) * pt.log(value))
+            + pt.switch(pt.eq(beta, 1.0), 0.0, (beta - 1.0) * pt.log1p(-value))
+            - (pt.gammaln(alpha) + pt.gammaln(beta) - pt.gammaln(alpha + beta))
+        )
+
+        # Cases
+        res = pt.switch(
+            pt.eq(value, 0),
+            pt.log(zoi) + pt.log1p(-coi),
+            pt.switch(
+                pt.eq(value, 1),
+                pt.log(zoi) + pt.log(coi),
+                pt.log1p(-zoi) + logp_beta,
+            ),
+        )
+
+        # Guard against values outside [0, 1]
+        res = pt.switch(
+            pt.bitwise_and(pt.ge(value, 0.0), pt.le(value, 1.0)),
+            res,
+            -np.inf,
+        )
+        return check_parameters(
+            res,
+            zoi >= 0,
+            zoi <= 1,
+            coi >= 0,
+            coi <= 1,
+            alpha > 0,
+            beta > 0,
+            msg="0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",
+        )
+
+
 class KumaraswamyRV(SymbolicRandomVariable):
     name = "kumaraswamy"
     extended_signature = "[rng],[size],(),()->[rng],()"
@@ -1318,13 +1554,23 @@ def dist(cls, a: DIST_PARAMETER_TYPES, b: DIST_PARAMETER_TYPES, *args, **kwargs)
         return super().dist([a, b], *args, **kwargs)
 
     def support_point(rv, size, a, b):
-        mean = pt.exp(pt.log(b) + pt.gammaln(1 + 1 / a) + pt.gammaln(b) - pt.gammaln(1 + 1 / a + b))
+        mean = pt.exp(
+            pt.log(b)
+            + pt.gammaln(1 + 1 / a)
+            + pt.gammaln(b)
+            - pt.gammaln(1 + 1 / a + b)
+        )
         if not rv_size_is_none(size):
             mean = pt.full(size, mean)
         return mean
 
     def logp(value, a, b):
-        res = pt.log(a) + pt.log(b) + (a - 1) * pt.log(value) + (b - 1) * pt.log(1 - value**a)
+        res = (
+            pt.log(a)
+            + pt.log(b)
+            + (a - 1) * pt.log(value)
+            + (b - 1) * pt.log(1 - value**a)
+        )
         res = pt.switch(
             pt.or_(pt.lt(value, 0), pt.gt(value, 1)),
             -np.inf,
@@ -1356,7 +1602,9 @@ def logcdf(value, a, b):
         )
 
     def icdf(value, a, b):
-        res = pt.exp(pt.reciprocal(a) * pt.log1mexp(pt.reciprocal(b) * pt.log1p(-value)))
+        res = pt.exp(
+            pt.reciprocal(a) * pt.log1mexp(pt.reciprocal(b) * pt.log1p(-value))
+        )
         res = check_icdf_value(res, value)
         return check_icdf_parameters(
             res,
@@ -1414,7 +1662,9 @@ def dist(cls, lam=None, *, scale=None, **kwargs):
         if lam is None and scale is None:
             scale = 1.0
         elif lam is not None and scale is not None:
-            raise ValueError("Incompatible parametrization. Can't specify both lam and scale.")
+            raise ValueError(
+                "Incompatible parametrization. Can't specify both lam and scale."
+            )
         elif lam is not None:
             scale = pt.reciprocal(lam)
 
@@ -1548,7 +1798,9 @@ def logcdf(value, mu, b):
 
     def icdf(value, mu, b):
         res = pt.switch(
-            pt.le(value, 0.5), mu + b * np.log(2 * value), mu - b * np.log(2 - 2 * value)
+            pt.le(value, 0.5),
+            mu + b * np.log(2 * value),
+            mu - b * np.log(2 - 2 * value),
         )
         res = check_icdf_value(res, value)
         return check_icdf_parameters(res, b > 0, msg="b > 0")
@@ -1568,7 +1820,9 @@ def rv_op(cls, b, kappa, mu, *, size=None, rng=None):
         size = normalize_size_param(size)
 
         if rv_size_is_none(size):
-            size = implicit_size_from_params(b, kappa, mu, ndims_params=cls.ndims_params)
+            size = implicit_size_from_params(
+                b, kappa, mu, ndims_params=cls.ndims_params
+            )
 
         next_rng, u = uniform(size=size, rng=rng, return_next_rng=True)
         switch = kappa**2 / (1 + kappa**2)
@@ -1921,8 +2175,10 @@ def logcdf(value, nu, mu, sigma):
     def icdf(value, nu, mu, sigma):
         res = pt.switch(
             pt.lt(value, 0.5),
-            -pt.sqrt(nu) * pt.sqrt((1.0 / betaincinv(nu * 0.5, 0.5, 2.0 * value)) - 1.0),
-            pt.sqrt(nu) * pt.sqrt((1.0 / betaincinv(nu * 0.5, 0.5, 2.0 * (1 - value))) - 1.0),
+            -pt.sqrt(nu)
+            * pt.sqrt((1.0 / betaincinv(nu * 0.5, 0.5, 2.0 * value)) - 1.0),
+            pt.sqrt(nu)
+            * pt.sqrt((1.0 / betaincinv(nu * 0.5, 0.5, 2.0 * (1 - value))) - 1.0),
         )
         res = mu + res * sigma
         res = check_icdf_value(res, value)
@@ -1943,7 +2199,9 @@ class SkewStudentTRV(RandomVariable):
     @classmethod
     def rng_fn(cls, rng, a, b, mu, sigma, size=None) -> np.ndarray:
         return np.asarray(
-            stats.jf_skew_t.rvs(a=a, b=b, loc=mu, scale=sigma, size=size, random_state=rng)
+            stats.jf_skew_t.rvs(
+                a=a, b=b, loc=mu, scale=sigma, size=size, random_state=rng
+            )
         )
 
 
@@ -2237,7 +2495,9 @@ def support_point(rv, size, alpha, beta):
         return alpha
 
     def logp(value, alpha, beta):
-        res = -pt.log(np.pi) - pt.log(beta) - pt.log1p(pt.pow((value - alpha) / beta, 2))
+        res = (
+            -pt.log(np.pi) - pt.log(beta) - pt.log1p(pt.pow((value - alpha) / beta, 2))
+        )
         return check_parameters(
             res,
             beta > 0,
@@ -2463,7 +2723,12 @@ def support_point(rv, size, alpha, scale):
 
     def logp(value, alpha, scale):
         beta = pt.reciprocal(scale)
-        res = -pt.gammaln(alpha) + logpow(beta, alpha) - beta * value + logpow(value, alpha - 1)
+        res = (
+            -pt.gammaln(alpha)
+            + logpow(beta, alpha)
+            - beta * value
+            + logpow(value, alpha - 1)
+        )
         res = pt.switch(pt.ge(value, 0.0), res, -np.inf)
         return check_parameters(
             res,
@@ -2581,7 +2846,12 @@ def _get_alpha_beta(cls, alpha, beta, mu, sigma):
         return alpha, beta
 
     def logp(value, alpha, beta):
-        res = -pt.gammaln(alpha) + logpow(beta, alpha) - beta / value + logpow(value, -alpha - 1)
+        res = (
+            -pt.gammaln(alpha)
+            + logpow(beta, alpha)
+            - beta / value
+            + logpow(value, -alpha - 1)
+        )
         res = pt.switch(pt.ge(value, 0.0), res, -np.inf)
         return check_parameters(
             res,
@@ -2808,7 +3078,9 @@ def rv_op(cls, nu, sigma, *, size=None, rng=None):
         next_rng, t_draws = t(df=nu, scale=sigma, size=size, rng=rng, return_next_rng=True)
         draws = pt.abs(t_draws)
 
-        return cls(inputs=[rng, size, nu, sigma], outputs=[next_rng, draws])(rng, size, nu, sigma)
+        return cls(inputs=[rng, size, nu, sigma], outputs=[next_rng, draws])(
+            rng, size, nu, sigma
+        )
 
 
 class HalfStudentT(PositiveContinuous):
@@ -2928,7 +3200,9 @@ def rv_op(cls, mu, sigma, nu, *, size=None, rng=None):
         size = normalize_size_param(size)
 
         if rv_size_is_none(size):
-            size = implicit_size_from_params(mu, sigma, nu, ndims_params=cls.ndims_params)
+            size = implicit_size_from_params(
+                mu, sigma, nu, ndims_params=cls.ndims_params
+            )
 
         next_rng, normal_draws = normal(
             loc=mu, scale=sigma, size=size, rng=rng, return_next_rng=True
@@ -3128,7 +3402,9 @@ def support_point(rv, size, mu, kappa):
 
     def logp(value, mu, kappa):
         res = kappa * pt.cos(mu - value) - pt.log(2 * np.pi) - pt.log(pt.i0(kappa))
-        res = pt.switch(pt.bitwise_and(pt.ge(value, -np.pi), pt.le(value, np.pi)), res, -np.inf)
+        res = pt.switch(
+            pt.bitwise_and(pt.ge(value, -np.pi), pt.le(value, np.pi)), res, -np.inf
+        )
         return check_parameters(
             res,
             kappa > 0,
@@ -3145,7 +3421,9 @@ class SkewNormalRV(RandomVariable):
     @classmethod
     def rng_fn(cls, rng, mu, sigma, alpha, size=None) -> np.ndarray:
         return np.asarray(
-            stats.skewnorm.rvs(a=alpha, loc=mu, scale=sigma, size=size, random_state=rng)
+            stats.skewnorm.rvs(
+                a=alpha, loc=mu, scale=sigma, size=size, random_state=rng
+            )
         )
 
 
@@ -3326,7 +3604,9 @@ def logp(value, lower, c, upper):
             pt.log(2 * (value - lower) / ((upper - lower) * (c - lower))),
             pt.log(2 * (upper - value) / ((upper - lower) * (upper - c))),
         )
-        res = pt.switch(pt.bitwise_and(pt.le(lower, value), pt.le(value, upper)), res, -np.inf)
+        res = pt.switch(
+            pt.bitwise_and(pt.le(lower, value), pt.le(value, upper)), res, -np.inf
+        )
         return check_parameters(
             res,
             lower <= c,
@@ -3762,7 +4042,11 @@ def dist(cls, mu=0, sigma=None, tau=None, **kwargs):
 
 
 def _interpolated_argcdf(p, pdf, cdf, x):
-    if np.prod(cdf.shape[:-1]) != 1 or np.prod(pdf.shape[:-1]) != 1 or np.prod(x.shape[:-1]) != 1:
+    if (
+        np.prod(cdf.shape[:-1]) != 1
+        or np.prod(pdf.shape[:-1]) != 1
+        or np.prod(x.shape[:-1]) != 1
+    ):
         raise NotImplementedError(
             "Function not implemented for batched points. "
             "Open an issue in https://github.com/pymc-devs/pymc if you need this functionality"
@@ -3782,7 +4066,9 @@ def _interpolated_argcdf(p, pdf, cdf, x):
     # This warning happens when we divide by slope = 0: we can ignore it
     # because the other result will be returned
     with warnings.catch_warnings():
-        warnings.filterwarnings("ignore", ".*invalid value encountered in.*", RuntimeWarning)
+        warnings.filterwarnings(
+            "ignore", ".*invalid value encountered in.*", RuntimeWarning
+        )
         large_slopes = (
             -pdf[index] + np.sqrt(pdf[index] ** 2 + 2 * slope * (p - cdf[index]))
         ) / slope
@@ -3880,7 +4166,8 @@ def support_point(rv, size, x_points, pdf_points, cdf_points):
         """Estimates the expectation integral using the trapezoid rule; cdf_points are not used."""
         x_fx = pt.mul(x_points, pdf_points)  # x_i * f(x_i) for all xi's in x_points
         support_point = (
-            pt.sum(pt.mul(pt.diff(x_points, axis=-1), x_fx[..., 1:] + x_fx[..., :-1])) / 2
+            pt.sum(pt.mul(pt.diff(x_points, axis=-1), x_fx[..., 1:] + x_fx[..., :-1]))
+            / 2
         )
 
         if not rv_size_is_none(size):
@@ -3891,7 +4178,9 @@ def support_point(rv, size, x_points, pdf_points, cdf_points):
     def logp(value, x_points, pdf_points, cdf_points):
         # x_points and pdf_points are expected to be non-symbolic arrays wrapped
         # within a tensor.constant. We use the .data method to retrieve them
-        interp = InterpolatedUnivariateSpline(x_points.data, pdf_points.data, k=1, ext="zeros")
+        interp = InterpolatedUnivariateSpline(
+            x_points.data, pdf_points.data, k=1, ext="zeros"
+        )
         Z = interp.integral(x_points.data[..., 0], x_points.data[..., -1])
 
         # interp and Z are converted to symbolic variables here
@@ -3991,7 +4280,11 @@ def support_point(rv, size, mu, sigma):
 
     def logp(value, mu, sigma):
         scaled = (value - mu) / sigma
-        res = -(1 / 2) * (scaled + pt.exp(-scaled)) - pt.log(sigma) - (1 / 2) * pt.log(2 * np.pi)
+        res = (
+            -(1 / 2) * (scaled + pt.exp(-scaled))
+            - pt.log(sigma)
+            - (1 / 2) * pt.log(2 * np.pi)
+        )
         return check_parameters(
             res,
             sigma > 0,
@@ -4052,7 +4345,9 @@ def rng_fn(cls, rng, h, z, size=None) -> np.ndarray:
         if size is None:
             size = np.broadcast_shapes(h.shape, z.shape)
         return np.asarray(
-            random_polyagamma(h, z, size=size, random_state=rng).astype(pytensor.config.floatX)
+            random_polyagamma(h, z, size=size, random_state=rng).astype(
+                pytensor.config.floatX
+            )
         )
 
 

From 3c1a8bb823966bf299da9c3f811996515c53c529 Mon Sep 17 00:00:00 2001
From: Noah Shin <noahshin0911@gmail.com>
Date: Sat, 9 May 2026 19:37:03 -1000
Subject: [PATCH 2/5] Adding ZeroOneInflatedBeta distribution

- logp, logcdf, support_point implementations
- mu/kappa parameterization with alternate alpha/beta parameterization
- ZeroOneInflatedBetaRV for sampling
- Full test suite
- API documentation updated
---
 docs/source/api/distributions/continuous.rst |   1 +
 pymc/distributions/continuous.py             |  34 +-
 tests/distributions/test_continuous.py       | 404 +++++++++++++++++--
 3 files changed, 398 insertions(+), 41 deletions(-)

diff --git a/docs/source/api/distributions/continuous.rst b/docs/source/api/distributions/continuous.rst
index 36b4070a92..1b12e7d6f8 100644
--- a/docs/source/api/distributions/continuous.rst
+++ b/docs/source/api/distributions/continuous.rst
@@ -41,3 +41,4 @@ Continuous
    VonMises
    Wald
    Weibull
+   ZeroOneInflatedBeta
diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py
index 77c552a865..97b8c82bee 100644
--- a/pymc/distributions/continuous.py
+++ b/pymc/distributions/continuous.py
@@ -1442,14 +1442,12 @@ def support_point(rv, size, zoi, coi, alpha, beta):
         return mean
 
     def logp(value, zoi, coi, alpha, beta):
-        # Beta log-probability
         logp_beta = (
             pt.switch(pt.eq(alpha, 1.0), 0.0, (alpha - 1.0) * pt.log(value))
             + pt.switch(pt.eq(beta, 1.0), 0.0, (beta - 1.0) * pt.log1p(-value))
             - (pt.gammaln(alpha) + pt.gammaln(beta) - pt.gammaln(alpha + beta))
         )
 
-        # Cases
         res = pt.switch(
             pt.eq(value, 0),
             pt.log(zoi) + pt.log1p(-coi),
@@ -1460,7 +1458,6 @@ def logp(value, zoi, coi, alpha, beta):
             ),
         )
 
-        # Guard against values outside [0, 1]
         res = pt.switch(
             pt.bitwise_and(pt.ge(value, 0.0), pt.le(value, 1.0)),
             res,
@@ -1474,7 +1471,36 @@ def logp(value, zoi, coi, alpha, beta):
             coi <= 1,
             alpha > 0,
             beta > 0,
-            msg="0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",
+            msg = "0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",
+        )
+    def logcdf(value, zoi, coi, alpha, beta):
+        beta_logcdf = pt.log(pt.betainc(alpha, beta, value))
+
+        res = pt.switch(
+            pt.lt(value, 0),
+            -np.inf,
+            pt.switch(
+                pt.eq(value, 0),
+                pt.log(zoi) + pt.log1p(-coi),
+                pt.switch(
+                    pt.lt(value, 1),
+                    pt.log(
+                        zoi * (1 - coi) + (1 - zoi) * pt.exp(beta_logcdf)
+                    ),
+                    0.0,
+                ),
+            ),
+        )
+
+        return check_parameters(
+            res,
+            zoi >= 0,
+            zoi <= 1,
+            coi >= 0,
+            coi <= 1,
+            alpha > 0,
+            beta > 0,
+            msg = "0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",
         )
 
 
diff --git a/tests/distributions/test_continuous.py b/tests/distributions/test_continuous.py
index 9f0c729e2c..17f3aebeb6 100644
--- a/tests/distributions/test_continuous.py
+++ b/tests/distributions/test_continuous.py
@@ -74,7 +74,9 @@ class TestBoundedContinuous:
     def get_dist_params_and_interval_bounds(self, model, rv_name):
         rv = model.named_vars[rv_name]
         dist_params = rv.owner.inputs
-        lower_interval, upper_interval = model.rvs_to_transforms[rv].args_fn(*rv.owner.inputs)
+        lower_interval, upper_interval = model.rvs_to_transforms[rv].args_fn(
+            *rv.owner.inputs
+        )
         return (
             dist_params,
             lower_interval,
@@ -197,20 +199,26 @@ def test_triangular(self):
             pm.Triangular,
             Runif,
             {"lower": -Rplusunif, "c": Runif, "upper": Rplusunif},
-            lambda value, c, lower, upper: st.triang.logpdf(value, c - lower, lower, upper - lower),
+            lambda value, c, lower, upper: st.triang.logpdf(
+                value, c - lower, lower, upper - lower
+            ),
             skip_paramdomain_outside_edge_test=True,
         )
         check_logcdf(
             pm.Triangular,
             Runif,
             {"lower": -Rplusunif, "c": Runif, "upper": Rplusunif},
-            lambda value, c, lower, upper: st.triang.logcdf(value, c - lower, lower, upper - lower),
+            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),
+            lambda q, c, lower, upper: st.triang.ppf(
+                q, c - lower, lower, upper - lower
+            ),
             skip_paramdomain_outside_edge_test=True,
         )
 
@@ -372,7 +380,9 @@ def test_wald_logp_custom_points(self, value, mu, lam, phi, alpha, logp):
         # See e.g., doi: 10.1111/j.1467-9876.2005.00510.x, or
         # http://www.gamlss.org/.
         with pm.Model() as model:
-            pm.Wald("wald", mu=mu, lam=lam, phi=phi, alpha=alpha, default_transform=None)
+            pm.Wald(
+                "wald", mu=mu, lam=lam, phi=phi, alpha=alpha, default_transform=None
+            )
         point = {"wald": value}
         decimals = select_by_precision(float64=6, float32=1)
         npt.assert_almost_equal(
@@ -420,10 +430,49 @@ def test_beta_icdf(self):
             lambda q, alpha, beta: st.beta.ppf(q, alpha, beta),
         )
 
+    def test_zero_one_inflated_beta_logp(self):
+        def zoib_logp(value, zoi, coi, alpha, beta):
+            if value == 0:
+                return np.log(zoi) + np.log1p(-coi)
+            elif value == 1:
+                return np.log(zoi) + np.log(coi)
+            else:
+                return np.log1p(-zoi) + st.beta.logpdf(value, alpha, beta)
+
+        def zoib_logcdf(value, zoi, coi, alpha, beta):
+            if value < 0:
+                return -np.inf
+            elif value == 0:
+                return np.log(zoi * (1 - coi))
+            elif value < 1:
+                return np.log(
+                    zoi * (1 - coi) + (1 - zoi) * st.beta.cdf(value, alpha, beta)
+                )
+            else:
+                return 0.0
+
+        check_logp(
+            pm.ZeroOneInflatedBeta,
+            Unit,
+            {"zoi": Runif, "coi": Runif, "alpha": Rplus, "beta": Rplus},
+            zoib_logp,
+        )
+        check_logcdf(
+            pm.ZeroOneInflatedBeta,
+            Unit,
+            {"zoi": Runif, "coi": Runif, "alpha": Rplus, "beta": Rplus},
+            zoib_logcdf,
+        )
+
     def test_kumaraswamy(self):
         # Scipy does not have a built-in Kumaraswamy
         def scipy_log_pdf(value, a, b):
-            return np.log(a) + np.log(b) + (a - 1) * np.log(value) + (b - 1) * np.log(1 - value**a)
+            return (
+                np.log(a)
+                + np.log(b)
+                + (a - 1) * np.log(value)
+                + (b - 1) * np.log(1 - value**a)
+            )
 
         def log1mexp(x):
             return np.log1p(-np.exp(x)) if x < np.log(0.5) else np.log(-np.expm1(x))
@@ -480,7 +529,9 @@ def test_laplace(self):
             {"mu": R, "b": Rplus},
             lambda value, mu, b: st.laplace.logcdf(value, mu, b),
         )
-        check_icdf(pm.Laplace, {"mu": R, "b": Rplus}, lambda q, mu, b: st.laplace.ppf(q, mu, b))
+        check_icdf(
+            pm.Laplace, {"mu": R, "b": Rplus}, lambda q, mu, b: st.laplace.ppf(q, mu, b)
+        )
 
     def test_laplace_asymmetric(self):
         check_logp(
@@ -496,13 +547,17 @@ def test_lognormal(self):
             pm.LogNormal,
             Rplus,
             {"mu": R, "tau": Rplusbig},
-            lambda value, mu, tau: floatX(st.lognorm.logpdf(value, tau**-0.5, 0, np.exp(mu))),
+            lambda value, mu, tau: floatX(
+                st.lognorm.logpdf(value, tau**-0.5, 0, np.exp(mu))
+            ),
         )
         check_logp(
             pm.LogNormal,
             Rplus,
             {"mu": R, "sigma": Rplusbig},
-            lambda value, mu, sigma: floatX(st.lognorm.logpdf(value, sigma, 0, np.exp(mu))),
+            lambda value, mu, sigma: floatX(
+                st.lognorm.logpdf(value, sigma, 0, np.exp(mu))
+            ),
         )
         check_logcdf(
             pm.LogNormal,
@@ -648,7 +703,9 @@ def test_half_cauchy(self):
             lambda value, beta: st.halfcauchy.logcdf(value, scale=beta),
         )
         check_icdf(
-            pm.HalfCauchy, {"beta": Rplusbig}, lambda q, beta: st.halfcauchy.ppf(q, scale=beta)
+            pm.HalfCauchy,
+            {"beta": Rplusbig},
+            lambda q, beta: st.halfcauchy.ppf(q, scale=beta),
         )
 
     def test_gamma_logp(self):
@@ -817,8 +874,20 @@ def test_skew_normal(self):
             (50.0, 50.000, 10.000, 10.000, -1.433714),
             (1000.0, 500.000, 10.000, 20.000, -1.573708e-11),
             (0.01, 0.01, 100.0, 0.01, -0.69314718),  # Fails in scipy version
-            (-0.43402407, 0.0, 0.1, 0.1, -13.59615423),  # Previous 32-bit version failed here
-            (-0.72402009, 0.0, 0.1, 0.1, -31.26571842),  # Previous 64-bit version failed here
+            (
+                -0.43402407,
+                0.0,
+                0.1,
+                0.1,
+                -13.59615423,
+            ),  # Previous 32-bit version failed here
+            (
+                -0.72402009,
+                0.0,
+                0.1,
+                0.1,
+                -31.26571842,
+            ),  # Previous 64-bit version failed here
         ],
     )
     def test_ex_gaussian_cdf(self, value, mu, sigma, nu, logcdf_val):
@@ -840,7 +909,9 @@ def test_ex_gaussian_cdf_outside_edges(self):
             skip_paramdomain_inside_edge_test=True,  # Valid values are tested above
         )
 
-    @pytest.mark.skipif(condition=(pytensor.config.floatX == "float32"), reason="Fails on float32")
+    @pytest.mark.skipif(
+        condition=(pytensor.config.floatX == "float32"), reason="Fails on float32"
+    )
     def test_vonmises(self):
         check_logp(
             pm.VonMises,
@@ -896,7 +967,8 @@ def test_logitnormal(self):
             Unit,
             {"mu": R, "sigma": Rplus},
             lambda value, mu, sigma: (
-                st.norm.logpdf(sp.logit(value), mu, sigma) - (np.log(value) + np.log1p(-value))
+                st.norm.logpdf(sp.logit(value), mu, sigma)
+                - (np.log(value) + np.log1p(-value))
             ),
             decimal=select_by_precision(float64=6, float32=1),
         )
@@ -919,14 +991,18 @@ def test_rice(self):
             lambda value, b, sigma: st.rice.logpdf(value, b=b, loc=0, scale=sigma),
         )
         if pytensor.config.floatX == "float32":
-            raise Exception("Flaky test: It passed this time, but XPASS is not allowed.")
+            raise Exception(
+                "Flaky test: It passed this time, but XPASS is not allowed."
+            )
 
     def test_rice_nu(self):
         check_logp(
             pm.Rice,
             Rplus,
             {"nu": Rplus, "sigma": Rplusbig},
-            lambda value, nu, sigma: st.rice.logpdf(value, b=nu / sigma, loc=0, scale=sigma),
+            lambda value, nu, sigma: st.rice.logpdf(
+                value, b=nu / sigma, loc=0, scale=sigma
+            ),
         )
 
     def test_moyal_logp(self):
@@ -954,7 +1030,9 @@ def test_moyal_logcdf(self):
             lambda value, mu, sigma: floatX(st.moyal.logcdf(value, mu, sigma)),
         )
         if pytensor.config.floatX == "float32":
-            raise Exception("Flaky test: It passed this time, but XPASS is not allowed.")
+            raise Exception(
+                "Flaky test: It passed this time, but XPASS is not allowed."
+            )
 
     def test_moyal_icdf(self):
         check_icdf(
@@ -976,7 +1054,9 @@ class TestedInterpolated(pm.Interpolated):
                     def dist(cls, **kwargs):
                         x_points = np.linspace(xmin, xmax, 100000)
                         pdf_points = st.norm.pdf(x_points, loc=mu, scale=sigma)
-                        return super().dist(x_points=x_points, pdf_points=pdf_points, **kwargs)
+                        return super().dist(
+                            x_points=x_points, pdf_points=pdf_points, **kwargs
+                        )
 
                 def ref_pdf(value):
                     return np.where(
@@ -1238,9 +1318,13 @@ def test_halfstudentt_support_point(self, nu, sigma, size, expected):
             (1, 1, [-np.inf, -np.inf, -np.inf], 10, None, np.full(3, 9)),
         ],
     )
-    def test_truncatednormal_support_point(self, mu, sigma, lower, upper, size, expected):
+    def test_truncatednormal_support_point(
+        self, mu, sigma, lower, upper, size, expected
+    ):
         with pm.Model() as model:
-            pm.TruncatedNormal("x", mu=mu, sigma=sigma, lower=lower, upper=upper, size=size)
+            pm.TruncatedNormal(
+                "x", mu=mu, sigma=sigma, lower=lower, upper=upper, size=size
+            )
         assert_support_point_is_expected(model, expected)
 
     @pytest.mark.parametrize(
@@ -1430,7 +1514,12 @@ def test_inverse_gamma_support_point(self, alpha, beta, size, expected):
         [
             (2, 1, None, 1 * 2 ** (1 / 2)),
             (2, 1, 5, np.full(5, 1 * 2 ** (1 / 2))),
-            (np.arange(2, 7), np.arange(1, 6), None, np.arange(1, 6) * 2 ** (1 / np.arange(2, 7))),
+            (
+                np.arange(2, 7),
+                np.arange(1, 6),
+                None,
+                np.arange(1, 6) * 2 ** (1 / np.arange(2, 7)),
+            ),
             (
                 np.arange(2, 7),
                 np.arange(1, 6),
@@ -1465,7 +1554,13 @@ def test_vonmises_support_point(self, mu, kappa, size, expected):
             (None, 1, 1, 5, np.full(5, 1)),
             (1, None, np.ones(5), None, np.full(5, 1)),
             (3, np.full(5, 2), None, None, np.full(5, 3)),
-            (np.arange(1, 6), None, np.arange(1, 6), (2, 5), np.full((2, 5), np.arange(1, 6))),
+            (
+                np.arange(1, 6),
+                None,
+                np.arange(1, 6),
+                (2, 5),
+                np.full((2, 5), np.arange(1, 6)),
+            ),
         ],
     )
     def test_wald_support_point(self, mu, lam, phi, size, expected):
@@ -1577,7 +1672,12 @@ def test_triangular_support_point(self, c, lower, upper, size, expected):
             (0, np.arange(1, 5), None, sp.expit(np.zeros(4))),
             (np.arange(4), 1, None, sp.expit(np.arange(4))),
             (1, 5, 4, sp.expit(np.ones(4))),
-            (np.arange(4), np.arange(1, 5), (2, 4), np.full((2, 4), sp.expit(np.arange(4)))),
+            (
+                np.arange(4),
+                np.arange(1, 5),
+                (2, 4),
+                np.full((2, 4), sp.expit(np.arange(4))),
+            ),
         ],
     )
     def test_logitnormal_support_point(self, mu, sigma, size, expected):
@@ -1633,7 +1733,12 @@ def test_interpolated_support_point(self, x_points, pdf_points, size, expected):
             (4.0, 3.0, None, 7.8110885363844345),
             (4.0, np.full(5, 3), None, np.full(5, 7.8110885363844345)),
             (np.arange(5), 1, None, np.arange(5) + 1.2703628454614782),
-            (np.arange(5), np.ones(5), (2, 5), np.full((2, 5), np.arange(5) + 1.2703628454614782)),
+            (
+                np.arange(5),
+                np.ones(5),
+                (2, 5),
+                np.full((2, 5), np.arange(5) + 1.2703628454614782),
+            ),
         ],
     )
     def test_moyal_support_point(self, mu, sigma, size, expected):
@@ -1871,7 +1976,9 @@ def asymmetriclaplace_rng_fn(self, b, kappa, mu, size, uniform_rng_fct):
 
     def seeded_asymmetriclaplace_rng_fn(self):
         uniform_rng_fct = self.get_random_state().uniform
-        return ft.partial(self.asymmetriclaplace_rng_fn, uniform_rng_fct=uniform_rng_fct)
+        return ft.partial(
+            self.asymmetriclaplace_rng_fn, uniform_rng_fct=uniform_rng_fct
+        )
 
     pymc_dist = pm.AsymmetricLaplace
 
@@ -1900,8 +2007,12 @@ class TestAsymmetricLaplaceQ(BaseTestDistributionRandom):
 
 
 class TestExGaussian(BaseTestDistributionRandom):
-    def exgaussian_rng_fn(self, mu, sigma, nu, size, normal_rng_fct, exponential_rng_fct):
-        return normal_rng_fct(mu, sigma, size=size) + exponential_rng_fct(scale=nu, size=size)
+    def exgaussian_rng_fn(
+        self, mu, sigma, nu, size, normal_rng_fct, exponential_rng_fct
+    ):
+        return normal_rng_fct(mu, sigma, size=size) + exponential_rng_fct(
+            scale=nu, size=size
+        )
 
     def seeded_exgaussian_rng_fn(self):
         normal_rng_fct = self.get_random_state().normal
@@ -1952,13 +2063,17 @@ class TestStudentT(BaseTestDistributionRandom):
 
 class TestHalfStudentT(BaseTestDistributionRandom):
     def halfstudentt_rng_fn(self, df, loc, scale, size, rng):
-        return np.abs(st.t.rvs(df=df, loc=loc, scale=scale, size=size, random_state=rng))
+        return np.abs(
+            st.t.rvs(df=df, loc=loc, scale=scale, size=size, random_state=rng)
+        )
 
     pymc_dist = pm.HalfStudentT
     pymc_dist_params = {"nu": 5.0, "sigma": 2.0}
     expected_rv_op_params = {"nu": 5.0, "sigma": 2.0}
     reference_dist_params = {"df": 5.0, "loc": 0, "scale": 2.0}
-    reference_dist = lambda self: ft.partial(self.halfstudentt_rng_fn, rng=self.get_random_state())  # noqa: E731
+    reference_dist = lambda self: ft.partial(
+        self.halfstudentt_rng_fn, rng=self.get_random_state()
+    )  # noqa: E731
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
@@ -2102,7 +2217,9 @@ class TestWald(BaseTestDistributionRandom):
 
     def check_pymc_draws_match_reference(self):
         npt.assert_array_almost_equal(
-            self.pymc_rv.eval(), self.reference_dist_draws + self.alpha, decimal=self.decimal
+            self.pymc_rv.eval(),
+            self.reference_dist_draws + self.alpha,
+            decimal=self.decimal,
         )
 
 
@@ -2192,7 +2309,9 @@ def logit_normal_rng_fn(self, rng, size, loc, scale):
     pymc_dist_params = {"mu": 5.0, "sigma": 10.0}
     expected_rv_op_params = {"mu": 5.0, "sigma": 10.0}
     reference_dist_params = {"loc": 5.0, "scale": 10.0}
-    reference_dist = lambda self: ft.partial(self.logit_normal_rng_fn, rng=self.get_random_state())  # noqa: E731
+    reference_dist = lambda self: ft.partial(
+        self.logit_normal_rng_fn, rng=self.get_random_state()
+    )  # noqa: E731
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
@@ -2263,7 +2382,9 @@ class TestBeta(BaseTestDistributionRandom):
     expected_rv_op_params = {"alpha": 2.0, "beta": 5.0}
     reference_dist_params = {"a": 2.0, "b": 5.0}
     size = 15
-    reference_dist = lambda self: ft.partial(clipped_beta_rvs, random_state=self.get_random_state())  # noqa: E731
+    reference_dist = lambda self: ft.partial(
+        clipped_beta_rvs, random_state=self.get_random_state()
+    )  # noqa: E731
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
@@ -2342,7 +2463,9 @@ def halfcauchy_rng_fn(self, scale, size, rng):
     pymc_dist_params = {"beta": 5.0}
     expected_rv_op_params = {"beta": 5.0}
     reference_dist_params = {"scale": 5.0}
-    reference_dist = lambda self: ft.partial(self.halfcauchy_rng_fn, rng=self.get_random_state())  # noqa: E731
+    reference_dist = lambda self: ft.partial(
+        self.halfcauchy_rng_fn, rng=self.get_random_state()
+    )  # noqa: E731
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference_not_numba",
@@ -2478,13 +2601,17 @@ class TestPolyaGamma(BaseTestDistributionRandom):
     def polyagamma_rng_fn(self, size, h, z, rng):
         # Polyagamma returns different values if inputs have explicit broadcasted dims
         # Which PyTensor RVs always do when size is not None.
-        return random_polyagamma(np.atleast_1d(h), np.atleast_1d(z), size=size, random_state=rng)
+        return random_polyagamma(
+            np.atleast_1d(h), np.atleast_1d(z), size=size, random_state=rng
+        )
 
     pymc_dist = pm.PolyaGamma
     pymc_dist_params = {"h": 1.0, "z": 0.0}
     expected_rv_op_params = {"h": 1.0, "z": 0.0}
     reference_dist_params = {"h": 1.0, "z": 0.0}
-    reference_dist = lambda self: ft.partial(self.polyagamma_rng_fn, rng=self.get_random_state())  # noqa: E731
+    reference_dist = lambda self: ft.partial(
+        self.polyagamma_rng_fn, rng=self.get_random_state()
+    )  # noqa: E731
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
@@ -2505,7 +2632,9 @@ def interpolated_rng_fn(self, size, mu, sigma, rng):
     pymc_dist_params = {"x_points": x_points, "pdf_points": pdf_points}
     reference_dist_params = {"mu": mu, "sigma": sigma}
 
-    reference_dist = lambda self: ft.partial(self.interpolated_rng_fn, rng=self.get_random_state())  # noqa: E731
+    reference_dist = lambda self: ft.partial(
+        self.interpolated_rng_fn, rng=self.get_random_state()
+    )  # noqa: E731
     checks_to_run = [
         "check_rv_size",
         "check_draws",
@@ -2526,7 +2655,9 @@ class TestedInterpolated(pm.Interpolated):
                     def dist(cls, **kwargs):
                         x_points = np.linspace(mu - 5 * sigma, mu + 5 * sigma, 100)
                         pdf_points = st.norm.pdf(x_points, loc=mu, scale=sigma)
-                        return super().dist(x_points=x_points, pdf_points=pdf_points, **kwargs)
+                        return super().dist(
+                            x_points=x_points, pdf_points=pdf_points, **kwargs
+                        )
 
                 continuous_random_tester(
                     TestedInterpolated,
@@ -2534,3 +2665,202 @@ def dist(cls, **kwargs):
                     extra_args={"rng": pytensor.shared(rng)},
                     ref_rand=ref_rand,
                 )
+
+
+class TestZeroOneInflatedBeta(BaseTestDistributionRandom):
+    pymc_dist = pm.ZeroOneInflatedBeta
+    pymc_dist_params = {"zoi": 0.3, "coi": 0.4, "mu": 0.5, "kappa": 10.0}
+    expected_rv_op_params = {"zoi": 0.3, "coi": 0.4, "alpha": 5.0, "beta_": 5.0}
+    size = 15
+    checks_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_rv_size",
+    ]
+
+
+class TestZeroOneInflatedBetaLogp:
+
+    def test_logp_at_zero(self):
+        zoi, coi, mu, kappa = 0.3, 0.4, 0.4, 10.0
+        expected = np.log(zoi) + np.log1p(-coi)
+        with pm.Model():
+            dist = pm.ZeroOneInflatedBeta.dist(zoi=zoi, coi=coi, mu=mu, kappa=kappa)
+            got = pm.logp(dist, 0.0).eval()
+        np.testing.assert_allclose(got, expected, rtol=1e-6)
+
+    def test_logp_at_one(self):
+        zoi, coi, mu, kappa = 0.3, 0.4, 0.4, 10.0
+        expected = np.log(zoi) + np.log(coi)
+        with pm.Model():
+            dist = pm.ZeroOneInflatedBeta.dist(zoi=zoi, coi=coi, mu=mu, kappa=kappa)
+            got = pm.logp(dist, 1.0).eval()
+        np.testing.assert_allclose(got, expected, rtol=1e-6)
+
+    def test_logp_interior(self):
+        zoi, coi, mu, kappa = 0.3, 0.4, 0.4, 10.0
+        alpha = mu * kappa
+        beta = (1 - mu) * kappa
+        for v in [0.1, 0.3, 0.5, 0.7, 0.9]:
+            expected = np.log1p(-zoi) + st.beta.logpdf(v, alpha, beta)
+            with pm.Model():
+                dist = pm.ZeroOneInflatedBeta.dist(zoi=zoi, coi=coi, mu=mu, kappa=kappa)
+                got = pm.logp(dist, v).eval()
+            np.testing.assert_allclose(
+                got, expected, rtol=1e-6, err_msg=f"Mismatch at value={v}"
+            )
+
+    def test_support_point(self):
+        zoi, coi, mu, kappa = 0.2, 0.3, 0.4, 10.0
+        expected = zoi * coi + (1 - zoi) * mu
+        with pm.Model():
+            dist = pm.ZeroOneInflatedBeta.dist(zoi=zoi, coi=coi, mu=mu, kappa=kappa)
+            alpha = mu * kappa
+            beta = (1 - mu) * kappa
+            sp = pm.distributions.continuous.ZeroOneInflatedBeta.support_point(
+                dist, None, zoi, coi, alpha, beta
+            )
+        np.testing.assert_allclose(float(sp), expected, rtol=1e-6)
+
+    def test_parameter_constraints(self):
+        from pymc.logprob.utils import ParameterValueError
+
+        with pytest.raises(ParameterValueError):
+            with pm.Model():
+                dist = pm.ZeroOneInflatedBeta.dist(zoi=1.5, coi=0.4, mu=0.5, kappa=10)
+                pm.logp(dist, 0.5).eval()
+        with pytest.raises(ParameterValueError):
+            with pm.Model():
+                dist = pm.ZeroOneInflatedBeta.dist(zoi=0.3, coi=1.5, mu=0.5, kappa=10)
+                pm.logp(dist, 0.5).eval()
+
+        with pytest.raises(ParameterValueError):
+            with pm.Model():
+                dist = pm.ZeroOneInflatedBeta.dist(
+                    zoi=0.3, coi=0.4, alpha=-1.0, beta=5.0
+                )
+                pm.logp(dist, 0.5).eval()
+
+        def test_zib_special_case(self):
+            with pm.Model():
+                y = pm.ZeroOneInflatedBeta(
+                    "y", zoi=0.5, coi=0.0, mu=0.5, kappa=10, size=1000
+                )
+                prior = pm.sample_prior_predictive(draws=1, random_seed=42)
+            samples = prior.prior["y"].values.flatten()
+            assert np.sum(samples == 1.0) == 0
+
+    def test_oib_special_case(self):
+        with pm.Model():
+            y = pm.ZeroOneInflatedBeta(
+                "y", zoi=0.5, coi=1.0, mu=0.5, kappa=10, size=1000
+            )
+            prior = pm.sample_prior_predictive(draws=1, random_seed=42)
+        samples = prior.prior["y"].values.flatten()
+        assert np.sum(samples == 0.0) == 0
+
+    def test_high_zoi_produces_many_boundaries(self):
+        with pm.Model():
+            y = pm.ZeroOneInflatedBeta(
+                "y", zoi=0.9, coi=0.5, mu=0.5, kappa=10, size=2000
+            )
+            prior = pm.sample_prior_predictive(draws=1, random_seed=42)
+        samples = prior.prior["y"].values.flatten()
+        prop_boundary = np.mean((samples == 0.0) | (samples == 1.0))
+        assert 0.83 < prop_boundary < 0.95
+
+    def test_zoib_both_boundaries(self):
+        with pm.Model():
+            y = pm.ZeroOneInflatedBeta(
+                "y", zoi=0.5, coi=0.5, mu=0.5, kappa=10, size=5000
+            )
+            prior = pm.sample_prior_predictive(draws=1, random_seed=42)
+        samples = prior.prior["y"].values.flatten()
+        prop_zeros = np.mean(samples == 0.0)
+        prop_ones = np.mean(samples == 1.0)
+        prop_interior = np.mean((samples > 0.0) & (samples < 1.0))
+        assert 0.18 < prop_zeros < 0.32
+        assert 0.18 < prop_ones < 0.32
+        assert 0.43 < prop_interior < 0.57
+
+    def test_zoib_reduces_to_beta(self):
+        mu, kappa = 0.4, 10.0
+        alpha = mu * kappa
+        beta = (1 - mu) * kappa
+        for v in [0.2, 0.5, 0.8]:
+            expected = st.beta.logpdf(v, alpha, beta)
+            with pm.Model():
+                dist = pm.ZeroOneInflatedBeta.dist(zoi=0.0, coi=0.5, mu=mu, kappa=kappa)
+                got = pm.logp(dist, v).eval()
+            np.testing.assert_allclose(
+                got, expected, rtol=1e-6, err_msg=f"Mismatch at value={v}"
+            )
+
+    def test_zoib_mu_kappa_matches_alpha_beta(self):
+        mu, kappa = 0.4, 10.0
+        alpha = mu * kappa
+        beta = (1 - mu) * kappa
+        zoi, coi = 0.3, 0.4
+        for v in [0.0, 0.5, 1.0]:
+            with pm.Model():
+                dist_mk = pm.ZeroOneInflatedBeta.dist(
+                    zoi=zoi, coi=coi, mu=mu, kappa=kappa
+                )
+                logp_mk = pm.logp(dist_mk, v).eval()
+            with pm.Model():
+                dist_ab = pm.ZeroOneInflatedBeta.dist(
+                    zoi=zoi, coi=coi, alpha=alpha, beta=beta
+                )
+                logp_ab = pm.logp(dist_ab, v).eval()
+            np.testing.assert_allclose(
+                logp_mk, logp_ab, rtol=1e-6, err_msg=f"Mismatch at value={v}"
+            )
+
+    def test_zoib_vectorized(self):
+        zoi = np.array([0.1, 0.2, 0.3])
+        coi = np.array([0.3, 0.4, 0.5])
+        mu = np.array([0.3, 0.5, 0.7])
+        kappa = np.array([10.0, 15.0, 20.0])
+        values = np.array([0.0, 0.5, 1.0])
+        with pm.Model():
+            dist = pm.ZeroOneInflatedBeta.dist(zoi=zoi, coi=coi, mu=mu, kappa=kappa)
+            logp = pm.logp(dist, values).eval()
+        assert logp.shape == (3,)
+        assert np.all(np.isfinite(logp))
+
+    def test_sampling_recovers_parameters(self):
+        rng = np.random.default_rng(42)
+        true_zoi, true_coi = 0.25, 0.40
+        true_mu, true_kappa = 0.45, 15.0
+        n = 600
+
+        alpha = true_mu * true_kappa
+        beta_p = (1 - true_mu) * true_kappa
+        u = rng.random(n)
+        data = np.where(
+            u < true_zoi * (1 - true_coi),
+            0.0,
+            np.where(u < true_zoi, 1.0, rng.beta(alpha, beta_p, n)),
+        )
+
+        with pm.Model():
+            zoi = pm.Beta("zoi", alpha=1, beta=1)
+            coi = pm.Beta("coi", alpha=1, beta=1)
+            mu = pm.Beta("mu", alpha=2, beta=2)
+            kappa = pm.Gamma("kappa", alpha=2, beta=0.1)
+            pm.ZeroOneInflatedBeta(
+                "y", zoi=zoi, coi=coi, mu=mu, kappa=kappa, observed=data
+            )
+            trace = pm.sample(
+                1000,
+                tune=1000,
+                chains=2,
+                random_seed=42,
+                progressbar=False,
+                return_inferencedata=True,
+            )
+
+        post = trace.posterior
+        assert abs(post["zoi"].mean().values - true_zoi) < 0.08
+        assert abs(post["coi"].mean().values - true_coi) < 0.08
+        assert abs(post["mu"].mean().values - true_mu) < 0.08
+        assert abs(post["kappa"].mean().values - true_kappa) < 5.0

From 93f14e86cac3f059a175251779a2bab055472a34 Mon Sep 17 00:00:00 2001
From: Noah Shin <noahshin0911@gmail.com>
Date: Sun, 10 May 2026 20:53:39 -1000
Subject: [PATCH 3/5] return_next_rng in lieu of .owner.outputs

---
 pymc/distributions/continuous.py | 166 ++++++++-----------------------
 1 file changed, 40 insertions(+), 126 deletions(-)

diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py
index 97b8c82bee..0b4c41d72f 100644
--- a/pymc/distributions/continuous.py
+++ b/pymc/distributions/continuous.py
@@ -174,9 +174,7 @@ def circ_cont_transform(op, rv):
 @_default_transform.register(BoundedContinuous)
 def bounded_cont_transform(op, rv, bound_args_indices=None):
     if bound_args_indices is None:
-        raise ValueError(
-            f"Must specify bound_args_indices for {op} bounded distribution"
-        )
+        raise ValueError(f"Must specify bound_args_indices for {op} bounded distribution")
 
     def transform_params(*args):
         lower, upper = None, None
@@ -446,9 +444,7 @@ def logp(value):
         return pt.switch(pt.lt(value, 0), -np.inf, pt.zeros_like(value))
 
     def logcdf(value):
-        return pt.switch(
-            pt.lt(value, np.inf), -np.inf, pt.switch(pt.eq(value, np.inf), 0, -np.inf)
-        )
+        return pt.switch(pt.lt(value, np.inf), -np.inf, pt.switch(pt.eq(value, np.inf), 0, -np.inf))
 
 
 class Normal(Continuous):
@@ -533,11 +529,7 @@ def support_point(rv, size, mu, sigma):
         return mu
 
     def logp(value, mu, sigma):
-        res = (
-            -0.5 * pt.pow((value - mu) / sigma, 2)
-            - pt.log(pt.sqrt(2.0 * np.pi))
-            - pt.log(sigma)
-        )
+        res = -0.5 * pt.pow((value - mu) / sigma, 2) - pt.log(pt.sqrt(2.0 * np.pi)) - pt.log(sigma)
         return check_parameters(
             res,
             sigma > 0,
@@ -694,12 +686,8 @@ def dist(
         sigma = pt.as_tensor_variable(sigma)
         mu = pt.as_tensor_variable(mu)
 
-        lower = (
-            pt.as_tensor_variable(lower) if lower is not None else pt.constant(-np.inf)
-        )
-        upper = (
-            pt.as_tensor_variable(upper) if upper is not None else pt.constant(np.inf)
-        )
+        lower = pt.as_tensor_variable(lower) if lower is not None else pt.constant(-np.inf)
+        upper = pt.as_tensor_variable(upper) if upper is not None else pt.constant(np.inf)
         return super().dist([mu, sigma, lower, upper], **kwargs)
 
     def support_point(rv, size, mu, sigma, lower, upper):
@@ -731,9 +719,7 @@ def logp(value, mu, sigma, lower, upper):
         is_lower_bounded = not (
             isinstance(lower, TensorConstant) and np.all(np.isneginf(lower.value))
         )
-        is_upper_bounded = not (
-            isinstance(upper, TensorConstant) and np.all(np.isinf(upper.value))
-        )
+        is_upper_bounded = not (isinstance(upper, TensorConstant) and np.all(np.isinf(upper.value)))
 
         if is_lower_bounded and is_upper_bounded:
             norm = log_diff_normal_cdf(mu, sigma, upper, lower)
@@ -769,9 +755,7 @@ def logcdf(value, mu, sigma, lower, upper):
         is_lower_bounded = not (
             isinstance(lower, TensorConstant) and np.all(np.isneginf(lower.value))
         )
-        is_upper_bounded = not (
-            isinstance(upper, TensorConstant) and np.all(np.isinf(upper.value))
-        )
+        is_upper_bounded = not (isinstance(upper, TensorConstant) and np.all(np.isinf(upper.value)))
 
         if is_lower_bounded:
             logcdf = pt.switch(value < lower, -np.inf, logcdf)
@@ -882,11 +866,7 @@ def support_point(rv, size, loc, sigma):
         return support_point
 
     def logp(value, loc, sigma):
-        res = (
-            -0.5 * pt.pow((value - loc) / sigma, 2)
-            + pt.log(pt.sqrt(2.0 / np.pi))
-            - pt.log(sigma)
-        )
+        res = -0.5 * pt.pow((value - loc) / sigma, 2) + pt.log(pt.sqrt(2.0 / np.pi)) - pt.log(sigma)
         res = pt.switch(pt.ge(value, loc), res, -np.inf)
         return check_parameters(
             res,
@@ -1232,9 +1212,7 @@ def logp(value, alpha, beta):
             + pt.switch(pt.eq(beta, 1.0), 0.0, (beta - 1.0) * pt.log1p(-value))
             - (pt.gammaln(alpha) + pt.gammaln(beta) - pt.gammaln(alpha + beta))
         )
-        res = pt.switch(
-            pt.bitwise_and(pt.ge(value, 0.0), pt.le(value, 1.0)), res, -np.inf
-        )
+        res = pt.switch(pt.bitwise_and(pt.ge(value, 0.0), pt.le(value, 1.0)), res, -np.inf)
         return check_parameters(
             res,
             alpha > 0,
@@ -1285,13 +1263,11 @@ def rv_op(cls, zoi, coi, alpha, beta_, *, rng=None, size=None):
         rng = normalize_rng_param(rng)
         size = normalize_size_param(size)
         if rv_size_is_none(size):
-            size = implicit_size_from_params(
-                zoi, coi, alpha, beta_, ndims_params=cls.ndims_params
-            )
+            size = implicit_size_from_params(zoi, coi, alpha, beta_, ndims_params=cls.ndims_params)
 
-        next_rng, u = uniform(size=size, rng=rng).owner.outputs
+        next_rng, u = uniform(size=size, rng=rng, return_next_rng=True)
 
-        next_rng, beta_draws = beta(alpha, beta_, size=size, rng=next_rng).owner.outputs
+        next_rng, beta_draws = beta(alpha, beta_, size=size, rng=next_rng, return_next_rng=True)
 
         draws = pt.switch(
             pt.lt(u, zoi * (1 - coi)),
@@ -1349,7 +1325,7 @@ class ZeroOneInflatedBeta(UnitContinuous):
         rng = np.random.default_rng(100)
         u = rng.random(n)
         alpha = true_mu * true_kappa
-        beta  = (1 - true_mu) * true_kappa   
+        beta  = (1 - true_mu) * true_kappa
         beta_data = rng.beta(alpha, beta, n)
         data = np.where(u < true_zoi * (1 - true_coi), 0.0,
                np.where(u < true_zoi, 1.0,
@@ -1411,11 +1387,7 @@ class ZeroOneInflatedBeta(UnitContinuous):
 
     @classmethod
     def dist(cls, zoi, coi, mu=None, kappa=None, alpha=None, beta=None, **kwargs):
-        if (
-            (alpha is not None)
-            and (beta is not None)
-            and (mu is not None or kappa is not None)
-        ):
+        if (alpha is not None) and (beta is not None) and (mu is not None or kappa is not None):
             warnings.warn(
                 "Both parametrizations provided. Proceeding with alpha and beta.",
                 UserWarning,
@@ -1471,8 +1443,9 @@ def logp(value, zoi, coi, alpha, beta):
             coi <= 1,
             alpha > 0,
             beta > 0,
-            msg = "0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",
+            msg="0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",
         )
+
     def logcdf(value, zoi, coi, alpha, beta):
         beta_logcdf = pt.log(pt.betainc(alpha, beta, value))
 
@@ -1484,9 +1457,7 @@ def logcdf(value, zoi, coi, alpha, beta):
                 pt.log(zoi) + pt.log1p(-coi),
                 pt.switch(
                     pt.lt(value, 1),
-                    pt.log(
-                        zoi * (1 - coi) + (1 - zoi) * pt.exp(beta_logcdf)
-                    ),
+                    pt.log(zoi * (1 - coi) + (1 - zoi) * pt.exp(beta_logcdf)),
                     0.0,
                 ),
             ),
@@ -1500,7 +1471,7 @@ def logcdf(value, zoi, coi, alpha, beta):
             coi <= 1,
             alpha > 0,
             beta > 0,
-            msg = "0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",
+            msg="0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",
         )
 
 
@@ -1580,23 +1551,13 @@ def dist(cls, a: DIST_PARAMETER_TYPES, b: DIST_PARAMETER_TYPES, *args, **kwargs)
         return super().dist([a, b], *args, **kwargs)
 
     def support_point(rv, size, a, b):
-        mean = pt.exp(
-            pt.log(b)
-            + pt.gammaln(1 + 1 / a)
-            + pt.gammaln(b)
-            - pt.gammaln(1 + 1 / a + b)
-        )
+        mean = pt.exp(pt.log(b) + pt.gammaln(1 + 1 / a) + pt.gammaln(b) - pt.gammaln(1 + 1 / a + b))
         if not rv_size_is_none(size):
             mean = pt.full(size, mean)
         return mean
 
     def logp(value, a, b):
-        res = (
-            pt.log(a)
-            + pt.log(b)
-            + (a - 1) * pt.log(value)
-            + (b - 1) * pt.log(1 - value**a)
-        )
+        res = pt.log(a) + pt.log(b) + (a - 1) * pt.log(value) + (b - 1) * pt.log(1 - value**a)
         res = pt.switch(
             pt.or_(pt.lt(value, 0), pt.gt(value, 1)),
             -np.inf,
@@ -1628,9 +1589,7 @@ def logcdf(value, a, b):
         )
 
     def icdf(value, a, b):
-        res = pt.exp(
-            pt.reciprocal(a) * pt.log1mexp(pt.reciprocal(b) * pt.log1p(-value))
-        )
+        res = pt.exp(pt.reciprocal(a) * pt.log1mexp(pt.reciprocal(b) * pt.log1p(-value)))
         res = check_icdf_value(res, value)
         return check_icdf_parameters(
             res,
@@ -1688,9 +1647,7 @@ def dist(cls, lam=None, *, scale=None, **kwargs):
         if lam is None and scale is None:
             scale = 1.0
         elif lam is not None and scale is not None:
-            raise ValueError(
-                "Incompatible parametrization. Can't specify both lam and scale."
-            )
+            raise ValueError("Incompatible parametrization. Can't specify both lam and scale.")
         elif lam is not None:
             scale = pt.reciprocal(lam)
 
@@ -1846,9 +1803,7 @@ def rv_op(cls, b, kappa, mu, *, size=None, rng=None):
         size = normalize_size_param(size)
 
         if rv_size_is_none(size):
-            size = implicit_size_from_params(
-                b, kappa, mu, ndims_params=cls.ndims_params
-            )
+            size = implicit_size_from_params(b, kappa, mu, ndims_params=cls.ndims_params)
 
         next_rng, u = uniform(size=size, rng=rng, return_next_rng=True)
         switch = kappa**2 / (1 + kappa**2)
@@ -2201,10 +2156,8 @@ def logcdf(value, nu, mu, sigma):
     def icdf(value, nu, mu, sigma):
         res = pt.switch(
             pt.lt(value, 0.5),
-            -pt.sqrt(nu)
-            * pt.sqrt((1.0 / betaincinv(nu * 0.5, 0.5, 2.0 * value)) - 1.0),
-            pt.sqrt(nu)
-            * pt.sqrt((1.0 / betaincinv(nu * 0.5, 0.5, 2.0 * (1 - value))) - 1.0),
+            -pt.sqrt(nu) * pt.sqrt((1.0 / betaincinv(nu * 0.5, 0.5, 2.0 * value)) - 1.0),
+            pt.sqrt(nu) * pt.sqrt((1.0 / betaincinv(nu * 0.5, 0.5, 2.0 * (1 - value))) - 1.0),
         )
         res = mu + res * sigma
         res = check_icdf_value(res, value)
@@ -2225,9 +2178,7 @@ class SkewStudentTRV(RandomVariable):
     @classmethod
     def rng_fn(cls, rng, a, b, mu, sigma, size=None) -> np.ndarray:
         return np.asarray(
-            stats.jf_skew_t.rvs(
-                a=a, b=b, loc=mu, scale=sigma, size=size, random_state=rng
-            )
+            stats.jf_skew_t.rvs(a=a, b=b, loc=mu, scale=sigma, size=size, random_state=rng)
         )
 
 
@@ -2521,9 +2472,7 @@ def support_point(rv, size, alpha, beta):
         return alpha
 
     def logp(value, alpha, beta):
-        res = (
-            -pt.log(np.pi) - pt.log(beta) - pt.log1p(pt.pow((value - alpha) / beta, 2))
-        )
+        res = -pt.log(np.pi) - pt.log(beta) - pt.log1p(pt.pow((value - alpha) / beta, 2))
         return check_parameters(
             res,
             beta > 0,
@@ -2749,12 +2698,7 @@ def support_point(rv, size, alpha, scale):
 
     def logp(value, alpha, scale):
         beta = pt.reciprocal(scale)
-        res = (
-            -pt.gammaln(alpha)
-            + logpow(beta, alpha)
-            - beta * value
-            + logpow(value, alpha - 1)
-        )
+        res = -pt.gammaln(alpha) + logpow(beta, alpha) - beta * value + logpow(value, alpha - 1)
         res = pt.switch(pt.ge(value, 0.0), res, -np.inf)
         return check_parameters(
             res,
@@ -2872,12 +2816,7 @@ def _get_alpha_beta(cls, alpha, beta, mu, sigma):
         return alpha, beta
 
     def logp(value, alpha, beta):
-        res = (
-            -pt.gammaln(alpha)
-            + logpow(beta, alpha)
-            - beta / value
-            + logpow(value, -alpha - 1)
-        )
+        res = -pt.gammaln(alpha) + logpow(beta, alpha) - beta / value + logpow(value, -alpha - 1)
         res = pt.switch(pt.ge(value, 0.0), res, -np.inf)
         return check_parameters(
             res,
@@ -3104,9 +3043,7 @@ def rv_op(cls, nu, sigma, *, size=None, rng=None):
         next_rng, t_draws = t(df=nu, scale=sigma, size=size, rng=rng, return_next_rng=True)
         draws = pt.abs(t_draws)
 
-        return cls(inputs=[rng, size, nu, sigma], outputs=[next_rng, draws])(
-            rng, size, nu, sigma
-        )
+        return cls(inputs=[rng, size, nu, sigma], outputs=[next_rng, draws])(rng, size, nu, sigma)
 
 
 class HalfStudentT(PositiveContinuous):
@@ -3226,9 +3163,7 @@ def rv_op(cls, mu, sigma, nu, *, size=None, rng=None):
         size = normalize_size_param(size)
 
         if rv_size_is_none(size):
-            size = implicit_size_from_params(
-                mu, sigma, nu, ndims_params=cls.ndims_params
-            )
+            size = implicit_size_from_params(mu, sigma, nu, ndims_params=cls.ndims_params)
 
         next_rng, normal_draws = normal(
             loc=mu, scale=sigma, size=size, rng=rng, return_next_rng=True
@@ -3428,9 +3363,7 @@ def support_point(rv, size, mu, kappa):
 
     def logp(value, mu, kappa):
         res = kappa * pt.cos(mu - value) - pt.log(2 * np.pi) - pt.log(pt.i0(kappa))
-        res = pt.switch(
-            pt.bitwise_and(pt.ge(value, -np.pi), pt.le(value, np.pi)), res, -np.inf
-        )
+        res = pt.switch(pt.bitwise_and(pt.ge(value, -np.pi), pt.le(value, np.pi)), res, -np.inf)
         return check_parameters(
             res,
             kappa > 0,
@@ -3447,9 +3380,7 @@ class SkewNormalRV(RandomVariable):
     @classmethod
     def rng_fn(cls, rng, mu, sigma, alpha, size=None) -> np.ndarray:
         return np.asarray(
-            stats.skewnorm.rvs(
-                a=alpha, loc=mu, scale=sigma, size=size, random_state=rng
-            )
+            stats.skewnorm.rvs(a=alpha, loc=mu, scale=sigma, size=size, random_state=rng)
         )
 
 
@@ -3630,9 +3561,7 @@ def logp(value, lower, c, upper):
             pt.log(2 * (value - lower) / ((upper - lower) * (c - lower))),
             pt.log(2 * (upper - value) / ((upper - lower) * (upper - c))),
         )
-        res = pt.switch(
-            pt.bitwise_and(pt.le(lower, value), pt.le(value, upper)), res, -np.inf
-        )
+        res = pt.switch(pt.bitwise_and(pt.le(lower, value), pt.le(value, upper)), res, -np.inf)
         return check_parameters(
             res,
             lower <= c,
@@ -4068,11 +3997,7 @@ def dist(cls, mu=0, sigma=None, tau=None, **kwargs):
 
 
 def _interpolated_argcdf(p, pdf, cdf, x):
-    if (
-        np.prod(cdf.shape[:-1]) != 1
-        or np.prod(pdf.shape[:-1]) != 1
-        or np.prod(x.shape[:-1]) != 1
-    ):
+    if np.prod(cdf.shape[:-1]) != 1 or np.prod(pdf.shape[:-1]) != 1 or np.prod(x.shape[:-1]) != 1:
         raise NotImplementedError(
             "Function not implemented for batched points. "
             "Open an issue in https://github.com/pymc-devs/pymc if you need this functionality"
@@ -4092,9 +4017,7 @@ def _interpolated_argcdf(p, pdf, cdf, x):
     # This warning happens when we divide by slope = 0: we can ignore it
     # because the other result will be returned
     with warnings.catch_warnings():
-        warnings.filterwarnings(
-            "ignore", ".*invalid value encountered in.*", RuntimeWarning
-        )
+        warnings.filterwarnings("ignore", ".*invalid value encountered in.*", RuntimeWarning)
         large_slopes = (
             -pdf[index] + np.sqrt(pdf[index] ** 2 + 2 * slope * (p - cdf[index]))
         ) / slope
@@ -4192,8 +4115,7 @@ def support_point(rv, size, x_points, pdf_points, cdf_points):
         """Estimates the expectation integral using the trapezoid rule; cdf_points are not used."""
         x_fx = pt.mul(x_points, pdf_points)  # x_i * f(x_i) for all xi's in x_points
         support_point = (
-            pt.sum(pt.mul(pt.diff(x_points, axis=-1), x_fx[..., 1:] + x_fx[..., :-1]))
-            / 2
+            pt.sum(pt.mul(pt.diff(x_points, axis=-1), x_fx[..., 1:] + x_fx[..., :-1])) / 2
         )
 
         if not rv_size_is_none(size):
@@ -4204,9 +4126,7 @@ def support_point(rv, size, x_points, pdf_points, cdf_points):
     def logp(value, x_points, pdf_points, cdf_points):
         # x_points and pdf_points are expected to be non-symbolic arrays wrapped
         # within a tensor.constant. We use the .data method to retrieve them
-        interp = InterpolatedUnivariateSpline(
-            x_points.data, pdf_points.data, k=1, ext="zeros"
-        )
+        interp = InterpolatedUnivariateSpline(x_points.data, pdf_points.data, k=1, ext="zeros")
         Z = interp.integral(x_points.data[..., 0], x_points.data[..., -1])
 
         # interp and Z are converted to symbolic variables here
@@ -4306,11 +4226,7 @@ def support_point(rv, size, mu, sigma):
 
     def logp(value, mu, sigma):
         scaled = (value - mu) / sigma
-        res = (
-            -(1 / 2) * (scaled + pt.exp(-scaled))
-            - pt.log(sigma)
-            - (1 / 2) * pt.log(2 * np.pi)
-        )
+        res = -(1 / 2) * (scaled + pt.exp(-scaled)) - pt.log(sigma) - (1 / 2) * pt.log(2 * np.pi)
         return check_parameters(
             res,
             sigma > 0,
@@ -4371,9 +4287,7 @@ def rng_fn(cls, rng, h, z, size=None) -> np.ndarray:
         if size is None:
             size = np.broadcast_shapes(h.shape, z.shape)
         return np.asarray(
-            random_polyagamma(h, z, size=size, random_state=rng).astype(
-                pytensor.config.floatX
-            )
+            random_polyagamma(h, z, size=size, random_state=rng).astype(pytensor.config.floatX)
         )
 
 

From 205d1497026bd4dd37a23f1a9e92ab5f7231f20a Mon Sep 17 00:00:00 2001
From: Noah Shin <noahshin0911@gmail.com>
Date: Sun, 10 May 2026 23:06:43 -1000
Subject: [PATCH 4/5] Rename alpha/beta_ to a/b in rv_op for consistency

---
 pymc/distributions/__init__.py         |   2 +-
 pymc/distributions/continuous.py       |  14 +--
 tests/distributions/test_continuous.py | 165 ++++++-------------------
 3 files changed, 49 insertions(+), 132 deletions(-)

diff --git a/pymc/distributions/__init__.py b/pymc/distributions/__init__.py
index 87494447eb..974100b230 100644
--- a/pymc/distributions/__init__.py
+++ b/pymc/distributions/__init__.py
@@ -51,7 +51,7 @@
     VonMises,
     Wald,
     Weibull,
-    ZeroOneInflatedBeta
+    ZeroOneInflatedBeta,
 )
 from pymc.distributions.custom import CustomDist, DensityDist
 from pymc.distributions.discrete import (
diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py
index 0b4c41d72f..5172e1b647 100644
--- a/pymc/distributions/continuous.py
+++ b/pymc/distributions/continuous.py
@@ -1255,19 +1255,19 @@ class ZeroOneInflatedBetaRV(SymbolicRandomVariable):
     _print_name = ("ZeroOneInflatedBeta", "\\operatorname{ZeroOneInflatedBeta}")
 
     @classmethod
-    def rv_op(cls, zoi, coi, alpha, beta_, *, rng=None, size=None):
+    def rv_op(cls, zoi, coi, a, b, *, rng=None, size=None):
         zoi = pt.as_tensor(zoi)
         coi = pt.as_tensor(coi)
-        alpha = pt.as_tensor(alpha)
-        beta_ = pt.as_tensor(beta_)
+        a = pt.as_tensor(a)
+        b = pt.as_tensor(b)
         rng = normalize_rng_param(rng)
         size = normalize_size_param(size)
         if rv_size_is_none(size):
-            size = implicit_size_from_params(zoi, coi, alpha, beta_, ndims_params=cls.ndims_params)
+            size = implicit_size_from_params(zoi, coi, a, b, ndims_params=cls.ndims_params)
 
         next_rng, u = uniform(size=size, rng=rng, return_next_rng=True)
 
-        next_rng, beta_draws = beta(alpha, beta_, size=size, rng=next_rng, return_next_rng=True)
+        next_rng, beta_draws = beta(a, b, size=size, rng=next_rng, return_next_rng=True)
 
         draws = pt.switch(
             pt.lt(u, zoi * (1 - coi)),
@@ -1280,9 +1280,9 @@ def rv_op(cls, zoi, coi, alpha, beta_, *, rng=None, size=None):
         )
 
         return cls(
-            inputs=[rng, size, zoi, coi, alpha, beta_],
+            inputs=[rng, size, zoi, coi, a, b],
             outputs=[next_rng, draws],
-        )(rng, size, zoi, coi, alpha, beta_)
+        )(rng, size, zoi, coi, a, b)
 
 
 class ZeroOneInflatedBeta(UnitContinuous):
diff --git a/tests/distributions/test_continuous.py b/tests/distributions/test_continuous.py
index 17f3aebeb6..9f584fd4d3 100644
--- a/tests/distributions/test_continuous.py
+++ b/tests/distributions/test_continuous.py
@@ -74,9 +74,7 @@ class TestBoundedContinuous:
     def get_dist_params_and_interval_bounds(self, model, rv_name):
         rv = model.named_vars[rv_name]
         dist_params = rv.owner.inputs
-        lower_interval, upper_interval = model.rvs_to_transforms[rv].args_fn(
-            *rv.owner.inputs
-        )
+        lower_interval, upper_interval = model.rvs_to_transforms[rv].args_fn(*rv.owner.inputs)
         return (
             dist_params,
             lower_interval,
@@ -199,26 +197,20 @@ def test_triangular(self):
             pm.Triangular,
             Runif,
             {"lower": -Rplusunif, "c": Runif, "upper": Rplusunif},
-            lambda value, c, lower, upper: st.triang.logpdf(
-                value, c - lower, lower, upper - lower
-            ),
+            lambda value, c, lower, upper: st.triang.logpdf(value, c - lower, lower, upper - lower),
             skip_paramdomain_outside_edge_test=True,
         )
         check_logcdf(
             pm.Triangular,
             Runif,
             {"lower": -Rplusunif, "c": Runif, "upper": Rplusunif},
-            lambda value, c, lower, upper: st.triang.logcdf(
-                value, c - lower, lower, upper - lower
-            ),
+            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
-            ),
+            lambda q, c, lower, upper: st.triang.ppf(q, c - lower, lower, upper - lower),
             skip_paramdomain_outside_edge_test=True,
         )
 
@@ -380,9 +372,7 @@ def test_wald_logp_custom_points(self, value, mu, lam, phi, alpha, logp):
         # See e.g., doi: 10.1111/j.1467-9876.2005.00510.x, or
         # http://www.gamlss.org/.
         with pm.Model() as model:
-            pm.Wald(
-                "wald", mu=mu, lam=lam, phi=phi, alpha=alpha, default_transform=None
-            )
+            pm.Wald("wald", mu=mu, lam=lam, phi=phi, alpha=alpha, default_transform=None)
         point = {"wald": value}
         decimals = select_by_precision(float64=6, float32=1)
         npt.assert_almost_equal(
@@ -445,9 +435,7 @@ def zoib_logcdf(value, zoi, coi, alpha, beta):
             elif value == 0:
                 return np.log(zoi * (1 - coi))
             elif value < 1:
-                return np.log(
-                    zoi * (1 - coi) + (1 - zoi) * st.beta.cdf(value, alpha, beta)
-                )
+                return np.log(zoi * (1 - coi) + (1 - zoi) * st.beta.cdf(value, alpha, beta))
             else:
                 return 0.0
 
@@ -467,12 +455,7 @@ def zoib_logcdf(value, zoi, coi, alpha, beta):
     def test_kumaraswamy(self):
         # Scipy does not have a built-in Kumaraswamy
         def scipy_log_pdf(value, a, b):
-            return (
-                np.log(a)
-                + np.log(b)
-                + (a - 1) * np.log(value)
-                + (b - 1) * np.log(1 - value**a)
-            )
+            return np.log(a) + np.log(b) + (a - 1) * np.log(value) + (b - 1) * np.log(1 - value**a)
 
         def log1mexp(x):
             return np.log1p(-np.exp(x)) if x < np.log(0.5) else np.log(-np.expm1(x))
@@ -529,9 +512,7 @@ def test_laplace(self):
             {"mu": R, "b": Rplus},
             lambda value, mu, b: st.laplace.logcdf(value, mu, b),
         )
-        check_icdf(
-            pm.Laplace, {"mu": R, "b": Rplus}, lambda q, mu, b: st.laplace.ppf(q, mu, b)
-        )
+        check_icdf(pm.Laplace, {"mu": R, "b": Rplus}, lambda q, mu, b: st.laplace.ppf(q, mu, b))
 
     def test_laplace_asymmetric(self):
         check_logp(
@@ -547,17 +528,13 @@ def test_lognormal(self):
             pm.LogNormal,
             Rplus,
             {"mu": R, "tau": Rplusbig},
-            lambda value, mu, tau: floatX(
-                st.lognorm.logpdf(value, tau**-0.5, 0, np.exp(mu))
-            ),
+            lambda value, mu, tau: floatX(st.lognorm.logpdf(value, tau**-0.5, 0, np.exp(mu))),
         )
         check_logp(
             pm.LogNormal,
             Rplus,
             {"mu": R, "sigma": Rplusbig},
-            lambda value, mu, sigma: floatX(
-                st.lognorm.logpdf(value, sigma, 0, np.exp(mu))
-            ),
+            lambda value, mu, sigma: floatX(st.lognorm.logpdf(value, sigma, 0, np.exp(mu))),
         )
         check_logcdf(
             pm.LogNormal,
@@ -909,9 +886,7 @@ def test_ex_gaussian_cdf_outside_edges(self):
             skip_paramdomain_inside_edge_test=True,  # Valid values are tested above
         )
 
-    @pytest.mark.skipif(
-        condition=(pytensor.config.floatX == "float32"), reason="Fails on float32"
-    )
+    @pytest.mark.skipif(condition=(pytensor.config.floatX == "float32"), reason="Fails on float32")
     def test_vonmises(self):
         check_logp(
             pm.VonMises,
@@ -967,8 +942,7 @@ def test_logitnormal(self):
             Unit,
             {"mu": R, "sigma": Rplus},
             lambda value, mu, sigma: (
-                st.norm.logpdf(sp.logit(value), mu, sigma)
-                - (np.log(value) + np.log1p(-value))
+                st.norm.logpdf(sp.logit(value), mu, sigma) - (np.log(value) + np.log1p(-value))
             ),
             decimal=select_by_precision(float64=6, float32=1),
         )
@@ -991,18 +965,14 @@ def test_rice(self):
             lambda value, b, sigma: st.rice.logpdf(value, b=b, loc=0, scale=sigma),
         )
         if pytensor.config.floatX == "float32":
-            raise Exception(
-                "Flaky test: It passed this time, but XPASS is not allowed."
-            )
+            raise Exception("Flaky test: It passed this time, but XPASS is not allowed.")
 
     def test_rice_nu(self):
         check_logp(
             pm.Rice,
             Rplus,
             {"nu": Rplus, "sigma": Rplusbig},
-            lambda value, nu, sigma: st.rice.logpdf(
-                value, b=nu / sigma, loc=0, scale=sigma
-            ),
+            lambda value, nu, sigma: st.rice.logpdf(value, b=nu / sigma, loc=0, scale=sigma),
         )
 
     def test_moyal_logp(self):
@@ -1030,9 +1000,7 @@ def test_moyal_logcdf(self):
             lambda value, mu, sigma: floatX(st.moyal.logcdf(value, mu, sigma)),
         )
         if pytensor.config.floatX == "float32":
-            raise Exception(
-                "Flaky test: It passed this time, but XPASS is not allowed."
-            )
+            raise Exception("Flaky test: It passed this time, but XPASS is not allowed.")
 
     def test_moyal_icdf(self):
         check_icdf(
@@ -1054,9 +1022,7 @@ class TestedInterpolated(pm.Interpolated):
                     def dist(cls, **kwargs):
                         x_points = np.linspace(xmin, xmax, 100000)
                         pdf_points = st.norm.pdf(x_points, loc=mu, scale=sigma)
-                        return super().dist(
-                            x_points=x_points, pdf_points=pdf_points, **kwargs
-                        )
+                        return super().dist(x_points=x_points, pdf_points=pdf_points, **kwargs)
 
                 def ref_pdf(value):
                     return np.where(
@@ -1318,13 +1284,9 @@ def test_halfstudentt_support_point(self, nu, sigma, size, expected):
             (1, 1, [-np.inf, -np.inf, -np.inf], 10, None, np.full(3, 9)),
         ],
     )
-    def test_truncatednormal_support_point(
-        self, mu, sigma, lower, upper, size, expected
-    ):
+    def test_truncatednormal_support_point(self, mu, sigma, lower, upper, size, expected):
         with pm.Model() as model:
-            pm.TruncatedNormal(
-                "x", mu=mu, sigma=sigma, lower=lower, upper=upper, size=size
-            )
+            pm.TruncatedNormal("x", mu=mu, sigma=sigma, lower=lower, upper=upper, size=size)
         assert_support_point_is_expected(model, expected)
 
     @pytest.mark.parametrize(
@@ -1976,9 +1938,7 @@ def asymmetriclaplace_rng_fn(self, b, kappa, mu, size, uniform_rng_fct):
 
     def seeded_asymmetriclaplace_rng_fn(self):
         uniform_rng_fct = self.get_random_state().uniform
-        return ft.partial(
-            self.asymmetriclaplace_rng_fn, uniform_rng_fct=uniform_rng_fct
-        )
+        return ft.partial(self.asymmetriclaplace_rng_fn, uniform_rng_fct=uniform_rng_fct)
 
     pymc_dist = pm.AsymmetricLaplace
 
@@ -2007,12 +1967,8 @@ class TestAsymmetricLaplaceQ(BaseTestDistributionRandom):
 
 
 class TestExGaussian(BaseTestDistributionRandom):
-    def exgaussian_rng_fn(
-        self, mu, sigma, nu, size, normal_rng_fct, exponential_rng_fct
-    ):
-        return normal_rng_fct(mu, sigma, size=size) + exponential_rng_fct(
-            scale=nu, size=size
-        )
+    def exgaussian_rng_fn(self, mu, sigma, nu, size, normal_rng_fct, exponential_rng_fct):
+        return normal_rng_fct(mu, sigma, size=size) + exponential_rng_fct(scale=nu, size=size)
 
     def seeded_exgaussian_rng_fn(self):
         normal_rng_fct = self.get_random_state().normal
@@ -2063,17 +2019,13 @@ class TestStudentT(BaseTestDistributionRandom):
 
 class TestHalfStudentT(BaseTestDistributionRandom):
     def halfstudentt_rng_fn(self, df, loc, scale, size, rng):
-        return np.abs(
-            st.t.rvs(df=df, loc=loc, scale=scale, size=size, random_state=rng)
-        )
+        return np.abs(st.t.rvs(df=df, loc=loc, scale=scale, size=size, random_state=rng))
 
     pymc_dist = pm.HalfStudentT
     pymc_dist_params = {"nu": 5.0, "sigma": 2.0}
     expected_rv_op_params = {"nu": 5.0, "sigma": 2.0}
     reference_dist_params = {"df": 5.0, "loc": 0, "scale": 2.0}
-    reference_dist = lambda self: ft.partial(
-        self.halfstudentt_rng_fn, rng=self.get_random_state()
-    )  # noqa: E731
+    reference_dist = lambda self: ft.partial(self.halfstudentt_rng_fn, rng=self.get_random_state())
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
@@ -2309,9 +2261,7 @@ def logit_normal_rng_fn(self, rng, size, loc, scale):
     pymc_dist_params = {"mu": 5.0, "sigma": 10.0}
     expected_rv_op_params = {"mu": 5.0, "sigma": 10.0}
     reference_dist_params = {"loc": 5.0, "scale": 10.0}
-    reference_dist = lambda self: ft.partial(
-        self.logit_normal_rng_fn, rng=self.get_random_state()
-    )  # noqa: E731
+    reference_dist = lambda self: ft.partial(self.logit_normal_rng_fn, rng=self.get_random_state())
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
@@ -2382,9 +2332,7 @@ class TestBeta(BaseTestDistributionRandom):
     expected_rv_op_params = {"alpha": 2.0, "beta": 5.0}
     reference_dist_params = {"a": 2.0, "b": 5.0}
     size = 15
-    reference_dist = lambda self: ft.partial(
-        clipped_beta_rvs, random_state=self.get_random_state()
-    )  # noqa: E731
+    reference_dist = lambda self: ft.partial(clipped_beta_rvs, random_state=self.get_random_state())
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
@@ -2463,9 +2411,7 @@ def halfcauchy_rng_fn(self, scale, size, rng):
     pymc_dist_params = {"beta": 5.0}
     expected_rv_op_params = {"beta": 5.0}
     reference_dist_params = {"scale": 5.0}
-    reference_dist = lambda self: ft.partial(
-        self.halfcauchy_rng_fn, rng=self.get_random_state()
-    )  # noqa: E731
+    reference_dist = lambda self: ft.partial(self.halfcauchy_rng_fn, rng=self.get_random_state())
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference_not_numba",
@@ -2601,17 +2547,13 @@ class TestPolyaGamma(BaseTestDistributionRandom):
     def polyagamma_rng_fn(self, size, h, z, rng):
         # Polyagamma returns different values if inputs have explicit broadcasted dims
         # Which PyTensor RVs always do when size is not None.
-        return random_polyagamma(
-            np.atleast_1d(h), np.atleast_1d(z), size=size, random_state=rng
-        )
+        return random_polyagamma(np.atleast_1d(h), np.atleast_1d(z), size=size, random_state=rng)
 
     pymc_dist = pm.PolyaGamma
     pymc_dist_params = {"h": 1.0, "z": 0.0}
     expected_rv_op_params = {"h": 1.0, "z": 0.0}
     reference_dist_params = {"h": 1.0, "z": 0.0}
-    reference_dist = lambda self: ft.partial(
-        self.polyagamma_rng_fn, rng=self.get_random_state()
-    )  # noqa: E731
+    reference_dist = lambda self: ft.partial(self.polyagamma_rng_fn, rng=self.get_random_state())
     checks_to_run = [
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
@@ -2632,9 +2574,7 @@ def interpolated_rng_fn(self, size, mu, sigma, rng):
     pymc_dist_params = {"x_points": x_points, "pdf_points": pdf_points}
     reference_dist_params = {"mu": mu, "sigma": sigma}
 
-    reference_dist = lambda self: ft.partial(
-        self.interpolated_rng_fn, rng=self.get_random_state()
-    )  # noqa: E731
+    reference_dist = lambda self: ft.partial(self.interpolated_rng_fn, rng=self.get_random_state())
     checks_to_run = [
         "check_rv_size",
         "check_draws",
@@ -2655,9 +2595,7 @@ class TestedInterpolated(pm.Interpolated):
                     def dist(cls, **kwargs):
                         x_points = np.linspace(mu - 5 * sigma, mu + 5 * sigma, 100)
                         pdf_points = st.norm.pdf(x_points, loc=mu, scale=sigma)
-                        return super().dist(
-                            x_points=x_points, pdf_points=pdf_points, **kwargs
-                        )
+                        return super().dist(x_points=x_points, pdf_points=pdf_points, **kwargs)
 
                 continuous_random_tester(
                     TestedInterpolated,
@@ -2670,7 +2608,7 @@ def dist(cls, **kwargs):
 class TestZeroOneInflatedBeta(BaseTestDistributionRandom):
     pymc_dist = pm.ZeroOneInflatedBeta
     pymc_dist_params = {"zoi": 0.3, "coi": 0.4, "mu": 0.5, "kappa": 10.0}
-    expected_rv_op_params = {"zoi": 0.3, "coi": 0.4, "alpha": 5.0, "beta_": 5.0}
+    expected_rv_op_params = {"zoi": 0.3, "coi": 0.4, "a": 5.0, "b": 5.0}
     size = 15
     checks_to_run = [
         "check_pymc_params_match_rv_op",
@@ -2679,7 +2617,6 @@ class TestZeroOneInflatedBeta(BaseTestDistributionRandom):
 
 
 class TestZeroOneInflatedBetaLogp:
-
     def test_logp_at_zero(self):
         zoi, coi, mu, kappa = 0.3, 0.4, 0.4, 10.0
         expected = np.log(zoi) + np.log1p(-coi)
@@ -2705,9 +2642,7 @@ def test_logp_interior(self):
             with pm.Model():
                 dist = pm.ZeroOneInflatedBeta.dist(zoi=zoi, coi=coi, mu=mu, kappa=kappa)
                 got = pm.logp(dist, v).eval()
-            np.testing.assert_allclose(
-                got, expected, rtol=1e-6, err_msg=f"Mismatch at value={v}"
-            )
+            np.testing.assert_allclose(got, expected, rtol=1e-6, err_msg=f"Mismatch at value={v}")
 
     def test_support_point(self):
         zoi, coi, mu, kappa = 0.2, 0.3, 0.4, 10.0
@@ -2735,34 +2670,26 @@ def test_parameter_constraints(self):
 
         with pytest.raises(ParameterValueError):
             with pm.Model():
-                dist = pm.ZeroOneInflatedBeta.dist(
-                    zoi=0.3, coi=0.4, alpha=-1.0, beta=5.0
-                )
+                dist = pm.ZeroOneInflatedBeta.dist(zoi=0.3, coi=0.4, alpha=-1.0, beta=5.0)
                 pm.logp(dist, 0.5).eval()
 
         def test_zib_special_case(self):
             with pm.Model():
-                y = pm.ZeroOneInflatedBeta(
-                    "y", zoi=0.5, coi=0.0, mu=0.5, kappa=10, size=1000
-                )
+                y = pm.ZeroOneInflatedBeta("y", zoi=0.5, coi=0.0, mu=0.5, kappa=10, size=1000)
                 prior = pm.sample_prior_predictive(draws=1, random_seed=42)
             samples = prior.prior["y"].values.flatten()
             assert np.sum(samples == 1.0) == 0
 
     def test_oib_special_case(self):
         with pm.Model():
-            y = pm.ZeroOneInflatedBeta(
-                "y", zoi=0.5, coi=1.0, mu=0.5, kappa=10, size=1000
-            )
+            y = pm.ZeroOneInflatedBeta("y", zoi=0.5, coi=1.0, mu=0.5, kappa=10, size=1000)
             prior = pm.sample_prior_predictive(draws=1, random_seed=42)
         samples = prior.prior["y"].values.flatten()
         assert np.sum(samples == 0.0) == 0
 
     def test_high_zoi_produces_many_boundaries(self):
         with pm.Model():
-            y = pm.ZeroOneInflatedBeta(
-                "y", zoi=0.9, coi=0.5, mu=0.5, kappa=10, size=2000
-            )
+            y = pm.ZeroOneInflatedBeta("y", zoi=0.9, coi=0.5, mu=0.5, kappa=10, size=2000)
             prior = pm.sample_prior_predictive(draws=1, random_seed=42)
         samples = prior.prior["y"].values.flatten()
         prop_boundary = np.mean((samples == 0.0) | (samples == 1.0))
@@ -2770,9 +2697,7 @@ def test_high_zoi_produces_many_boundaries(self):
 
     def test_zoib_both_boundaries(self):
         with pm.Model():
-            y = pm.ZeroOneInflatedBeta(
-                "y", zoi=0.5, coi=0.5, mu=0.5, kappa=10, size=5000
-            )
+            y = pm.ZeroOneInflatedBeta("y", zoi=0.5, coi=0.5, mu=0.5, kappa=10, size=5000)
             prior = pm.sample_prior_predictive(draws=1, random_seed=42)
         samples = prior.prior["y"].values.flatten()
         prop_zeros = np.mean(samples == 0.0)
@@ -2791,9 +2716,7 @@ def test_zoib_reduces_to_beta(self):
             with pm.Model():
                 dist = pm.ZeroOneInflatedBeta.dist(zoi=0.0, coi=0.5, mu=mu, kappa=kappa)
                 got = pm.logp(dist, v).eval()
-            np.testing.assert_allclose(
-                got, expected, rtol=1e-6, err_msg=f"Mismatch at value={v}"
-            )
+            np.testing.assert_allclose(got, expected, rtol=1e-6, err_msg=f"Mismatch at value={v}")
 
     def test_zoib_mu_kappa_matches_alpha_beta(self):
         mu, kappa = 0.4, 10.0
@@ -2802,14 +2725,10 @@ def test_zoib_mu_kappa_matches_alpha_beta(self):
         zoi, coi = 0.3, 0.4
         for v in [0.0, 0.5, 1.0]:
             with pm.Model():
-                dist_mk = pm.ZeroOneInflatedBeta.dist(
-                    zoi=zoi, coi=coi, mu=mu, kappa=kappa
-                )
+                dist_mk = pm.ZeroOneInflatedBeta.dist(zoi=zoi, coi=coi, mu=mu, kappa=kappa)
                 logp_mk = pm.logp(dist_mk, v).eval()
             with pm.Model():
-                dist_ab = pm.ZeroOneInflatedBeta.dist(
-                    zoi=zoi, coi=coi, alpha=alpha, beta=beta
-                )
+                dist_ab = pm.ZeroOneInflatedBeta.dist(zoi=zoi, coi=coi, alpha=alpha, beta=beta)
                 logp_ab = pm.logp(dist_ab, v).eval()
             np.testing.assert_allclose(
                 logp_mk, logp_ab, rtol=1e-6, err_msg=f"Mismatch at value={v}"
@@ -2847,9 +2766,7 @@ def test_sampling_recovers_parameters(self):
             coi = pm.Beta("coi", alpha=1, beta=1)
             mu = pm.Beta("mu", alpha=2, beta=2)
             kappa = pm.Gamma("kappa", alpha=2, beta=0.1)
-            pm.ZeroOneInflatedBeta(
-                "y", zoi=zoi, coi=coi, mu=mu, kappa=kappa, observed=data
-            )
+            pm.ZeroOneInflatedBeta("y", zoi=zoi, coi=coi, mu=mu, kappa=kappa, observed=data)
             trace = pm.sample(
                 1000,
                 tune=1000,

From 88559d2c5695fcd9820db0e3815b6d104586ff42 Mon Sep 17 00:00:00 2001
From: Noah Shin <noahshin0911@gmail.com>
Date: Sat, 16 May 2026 14:59:25 -1000
Subject: [PATCH 5/5] Adding guards in logp and logcdf for stability via
 pt.clip

---
 pymc/distributions/continuous.py | 38 +++++++++++++++++++++++---------
 1 file changed, 28 insertions(+), 10 deletions(-)

diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py
index 5172e1b647..7b1a806aca 100644
--- a/pymc/distributions/continuous.py
+++ b/pymc/distributions/continuous.py
@@ -1414,9 +1414,19 @@ def support_point(rv, size, zoi, coi, alpha, beta):
         return mean
 
     def logp(value, zoi, coi, alpha, beta):
+        saved_zoi = zoi
+        saved_coi = coi    
+
+        check_zoi = pt.bitwise_and(pt.ge(zoi, 0), pt.le(zoi, 1))
+        check_coi = pt.bitwise_and(pt.ge(coi, 0), pt.le(coi, 1))
+        zoi = pt.switch(check_zoi, pt.clip(zoi, 1e-12, 1 - 1e-12), zoi)
+        coi = pt.switch(check_coi, pt.clip(coi, 1e-12, 1 - 1e-12), coi)
+
+        safe_value = pt.clip(value, 1e-12, 1 - 1e-12)
+
         logp_beta = (
-            pt.switch(pt.eq(alpha, 1.0), 0.0, (alpha - 1.0) * pt.log(value))
-            + pt.switch(pt.eq(beta, 1.0), 0.0, (beta - 1.0) * pt.log1p(-value))
+            pt.switch(pt.eq(alpha, 1.0), 0.0, (alpha - 1.0) * pt.log(safe_value))
+            + pt.switch(pt.eq(beta, 1.0), 0.0, (beta - 1.0) * pt.log1p(-safe_value))
             - (pt.gammaln(alpha) + pt.gammaln(beta) - pt.gammaln(alpha + beta))
         )
 
@@ -1437,16 +1447,24 @@ def logp(value, zoi, coi, alpha, beta):
         )
         return check_parameters(
             res,
-            zoi >= 0,
-            zoi <= 1,
-            coi >= 0,
-            coi <= 1,
+            saved_zoi >= 0,
+            saved_zoi <= 1,
+            saved_coi >= 0,
+            saved_coi <= 1,
             alpha > 0,
             beta > 0,
             msg="0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",
         )
 
     def logcdf(value, zoi, coi, alpha, beta):
+        saved_zoi = zoi
+        saved_coi = coi
+
+        check_zoi = pt.bitwise_and(pt.ge(zoi, 0), pt.le(zoi, 1))
+        check_coi = pt.bitwise_and(pt.ge(coi, 0), pt.le(coi, 1))
+        zoi = pt.switch(check_zoi, pt.clip(zoi, 1e-12, 1 - 1e-12), zoi)
+        coi = pt.switch(check_coi, pt.clip(coi, 1e-12, 1 - 1e-12), coi)
+
         beta_logcdf = pt.log(pt.betainc(alpha, beta, value))
 
         res = pt.switch(
@@ -1465,10 +1483,10 @@ def logcdf(value, zoi, coi, alpha, beta):
 
         return check_parameters(
             res,
-            zoi >= 0,
-            zoi <= 1,
-            coi >= 0,
-            coi <= 1,
+            saved_zoi >= 0,
+            saved_zoi <= 1,
+            saved_coi >= 0,
+            saved_coi <= 1,
             alpha > 0,
             beta > 0,
             msg="0 <= zoi <= 1, 0 <= coi <= 1, alpha > 0, beta > 0",