Use automatic logprob for LogitNormal and enable icdf

Co-authored-by: Ricardo Vieira <28983449+ricardov94@users.noreply.github.com>

Author: jnetzel1

pymc/distributions/continuous.py

@@ -53,6 +53,7 @@
 from pytensor.tensor.random.utils import normalize_size_param
 from pytensor.tensor.variable import TensorConstant, TensorVariable
 
+from pymc.distributions.custom import CustomDist
 from pymc.logprob.abstract import _logprob_helper
 from pymc.logprob.basic import TensorLike, icdf
 from pymc.pytensorf import normalize_rng_param
@@ -92,7 +93,7 @@ def polyagamma_cdf(*args, **kwargs):
 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, logit
+from pymc.math import invlogit, logdiffexp
 
 __all__ = [
     "AsymmetricLaplace",
@@ -3603,28 +3604,7 @@ def icdf(value, mu, s):
         )
 
 
-class LogitNormalRV(SymbolicRandomVariable):
-    name = "logit_normal"
-    extended_signature = "[rng],[size],(),()->[rng],()"
-    _print_name = ("LogitNormal", "\\operatorname{LogitNormal}")
-
-    @classmethod
-    def rv_op(cls, mu, sigma, *, size=None, rng=None):
-        mu = pt.as_tensor(mu)
-        sigma = pt.as_tensor(sigma)
-        rng = normalize_rng_param(rng)
-        size = normalize_size_param(size)
-
-        next_rng, normal_draws = normal(loc=mu, scale=sigma, size=size, rng=rng).owner.outputs
-        draws = pt.expit(normal_draws)
-
-        return cls(
-            inputs=[rng, size, mu, sigma],
-            outputs=[next_rng, draws],
-        )(rng, size, mu, sigma)
-
-
-class LogitNormal(UnitContinuous):
+class LogitNormal:
     r"""
     Logit-Normal distribution.
 
@@ -3672,37 +3652,26 @@ class LogitNormal(UnitContinuous):
         Defaults to 1.
     """
 
-    rv_type = LogitNormalRV
-    rv_op = LogitNormalRV.rv_op
+    @staticmethod
+    def logitnormal_dist(mu, sigma, size):
+        return invlogit(Normal.dist(mu=mu, sigma=sigma, size=size))
 
-    @classmethod
-    def dist(cls, mu=0, sigma=None, tau=None, **kwargs):
+    def __new__(cls, name, mu=0, sigma=None, tau=None, **kwargs):
         _, sigma = get_tau_sigma(tau=tau, sigma=sigma)
-        return super().dist([mu, sigma], **kwargs)
-
-    def support_point(rv, size, mu, sigma):
-        median, _ = pt.broadcast_arrays(invlogit(mu), sigma)
-        if not rv_size_is_none(size):
-            median = pt.full(size, median)
-        return median
-
-    def logp(value, mu, sigma):
-        tau, _ = get_tau_sigma(sigma=sigma)
-
-        res = pt.switch(
-            pt.or_(pt.le(value, 0), pt.ge(value, 1)),
-            -np.inf,
-            (
-                -0.5 * tau * (logit(value) - mu) ** 2
-                + 0.5 * pt.log(tau / (2.0 * np.pi))
-                - pt.log(value * (1 - value))
-            ),
+        return CustomDist(
+            name,
+            mu,
+            sigma,
+            dist=cls.logitnormal_dist,
+            class_name="LogitNormal",
+            **kwargs,
         )
 
-        return check_parameters(
-            res,
-            tau > 0,
-            msg="tau > 0",
+    @classmethod
+    def dist(cls, mu=0, sigma=None, tau=None, **kwargs):
+        _, sigma = get_tau_sigma(tau=tau, sigma=sigma)
+        return CustomDist.dist(
+            mu, sigma, dist=cls.logitnormal_dist, class_name="LogitNormal", **kwargs
         )
 
 

pymc/logprob/transforms.py

@@ -178,6 +178,7 @@ class MeasurableTransform(MeasurableElemwise):
         Erf,
         Erfc,
         Erfcx,
+        Sigmoid,
     )
 
     # Cannot use `transform` as name because it would clash with the property added by
@@ -227,7 +228,7 @@ def measurable_transform_logprob(op: MeasurableTransform, values, *inputs, **kwa
     return pt.switch(pt.isnan(jacobian), -np.inf, input_logprob + jacobian)
 
 
-MONOTONICALLY_INCREASING_OPS = (Exp, Log, Add, Sinh, Tanh, ArcSinh, ArcCosh, ArcTanh, Erf)
+MONOTONICALLY_INCREASING_OPS = (Exp, Log, Add, Sinh, Tanh, ArcSinh, ArcCosh, ArcTanh, Erf, Sigmoid)
 MONOTONICALLY_DECREASING_OPS = (Erfc, Erfcx)
 
 
@@ -401,7 +402,7 @@ def measurable_special_log_to_log(fgraph, node):
         return [pt.log(inp) / pt.log(10)]
 
 
-@node_rewriter([expm1, sigmoid, exp2])
+@node_rewriter([expm1, exp2])
 def measurable_special_exp_to_exp(fgraph, node):
     """Convert expm1, sigmoid, and exp2 of `MeasurableVariable`s to xp form."""
     if not filter_measurable_variables(node.inputs):
@@ -412,8 +413,6 @@ def measurable_special_exp_to_exp(fgraph, node):
         return [pt.exp(pt.log(2) * inp)]
     if isinstance(node.op.scalar_op, Expm1):
         return [pt.add(pt.exp(inp), -1)]
-    if isinstance(node.op.scalar_op, Sigmoid):
-        return [1 / (1 + pt.exp(-inp))]
 
 
 @node_rewriter([pow])
@@ -451,6 +450,7 @@ def measurable_power_exponent_to_exp(fgraph, node):
         erf,
         erfc,
         erfcx,
+        sigmoid,
     ]
 )
 def find_measurable_transforms(fgraph: FunctionGraph, node: Apply) -> list[Variable] | None:
@@ -526,6 +526,7 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Apply) -> list[Varia
             Erf: ErfTransform,
             Erfc: ErfcTransform,
             Erfcx: ErfcxTransform,
+            Sigmoid: SigmoidTransform,
         }[type(scalar_op)]()
 
     transform_op = MeasurableTransform(
@@ -968,6 +969,19 @@ def log_jac_det(self, value, *inputs):
         return pt.log(sigmoid_value) + pt.log1p(-sigmoid_value)
 
 
+class SigmoidTransform(Transform):
+    name = "sigmoid"
+
+    def forward(self, value, *inputs):
+        return sigmoid(value)  # AKA invlogit/expit
+
+    def backward(self, value, *inputs):
+        return pt.log(value) - pt.log1p(-value)
+
+    def log_jac_det(self, value, *inputs):
+        return -pt.log(value) - pt.log1p(-value)
+
+
 class SimplexTransform(Transform):
     name = "simplex"
 

tests/distributions/test_continuous.py

@@ -872,6 +872,12 @@ def test_logitnormal(self):
             ),
             decimal=select_by_precision(float64=6, float32=1),
         )
+        check_icdf(
+            pm.LogitNormal,
+            {"mu": R, "sigma": Rplus},
+            lambda q, mu, sigma: sp.expit(mu + sigma * st.norm.ppf(q)),
+            decimal=select_by_precision(float64=12, float32=5),
+        )
 
     @pytest.mark.skipif(
         condition=(pytensor.config.floatX == "float32"),