doc(gh-2187): docstrings for Gumbel, Laplace, Logistic (#2202)

* doc(gh-2187): enhance documentation for Gumbel, Laplace, and Logistic distributions

Adds math-flavored docstrings to three more continuous distributions on the
#2187 checklist, following the template from #2188, #2192, and #2199:

- Gumbel: class, __init__, sample, log_prob, mean, variance, cdf, icdf.
- Laplace: class, __init__, sample, log_prob, mean, variance, cdf, icdf, entropy.
- Logistic: class, __init__, sample, log_prob, mean, variance, cdf, icdf, entropy.

Docstrings only, no logic changes. ruff check and ruff format --check are
clean.

* fix: small fixes

---------

Co-authored-by: Meesum Qazalbash <meesumqazalbash@gmail.com>

Author: Metbcy

numpyro/distributions/continuous.py

@@ -841,7 +841,7 @@ def variance(self) -> ArrayLike:
         r"""If :math:`X \sim \mathrm{Gamma}(\alpha, \lambda)`, then
 
         .. math::
-            \mathrm{Var}[X] = \frac{\alpha}{\lambda^2}
+            \mathrm{Var}(X) = \frac{\alpha}{\lambda^2}
         """
         return self.concentration / jnp.power(self.rate, 2)
 
@@ -1306,6 +1306,22 @@ def mean(self) -> ArrayLike:
 
 
 class Gumbel(Distribution):
+    r"""The Gumbel (maximum) distribution, a continuous real-valued
+    distribution parameterized by location :math:`\mu` and scale :math:`\beta > 0`.
+    It is the limiting distribution of the maximum of a large number of i.i.d.
+    samples from an exponential-tailed distribution.
+
+    The Probability Density Function (PDF) is:
+
+    .. math::
+        f(x \mid \mu, \beta) = \frac{1}{\beta} \exp\!\left(
+            -\frac{x - \mu}{\beta} - \exp\!\left(-\frac{x - \mu}{\beta}\right)
+        \right), \quad x \in \mathbb{R}
+
+    where :math:`\mu \in \mathbb{R}` is the location (:attr:`loc`) and
+    :math:`\beta > 0` is the scale (:attr:`scale`).
+    """
+
     arg_constraints = {"loc": constraints.real, "scale": constraints.positive}
     support = constraints.real
     reparametrized_params = ["loc", "scale"]
@@ -1317,6 +1333,11 @@ def __init__(
         *,
         validate_args: Optional[bool] = None,
     ) -> None:
+        r"""
+        :param loc: Location parameter :math:`\mu \in \mathbb{R}`. Defaults to ``0.0``.
+        :param scale: Scale parameter :math:`\beta > 0`. Defaults to ``1.0``.
+        :param validate_args: If True, enforce domain constraints during initialization.
+        """
         self.loc, self.scale = promote_shapes(loc, scale)
         batch_shape = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
 
@@ -1325,6 +1346,15 @@ def __init__(
         )
 
     def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+        r"""Draw samples from the Gumbel distribution via the location-scale
+        transform :math:`X = \mu + \beta Z`, where
+        :math:`Z \sim \mathrm{Gumbel}(0, 1)` is drawn from
+        :func:`~jax.random.gumbel`.
+
+        :param key: A JAX PRNG key.
+        :param sample_shape: Sample dimensions to prepend to the batch shape.
+        :return: Real-valued samples from the Gumbel distribution.
+        """
         assert is_prng_key(key)
         standard_gumbel_sample = random.gumbel(
             key, shape=sample_shape + self.batch_shape + self.event_shape
@@ -1333,23 +1363,63 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
 
     @validate_sample
     def log_prob(self, value: ArrayLike) -> ArrayLike:
+        r"""Evaluate the log probability density function at ``value``.
+
+        Letting :math:`z = (x - \mu)/\beta`,
+
+        .. math::
+            \ln f(x \mid \mu, \beta) = -z - e^{-z} - \ln \beta
+
+        :param value: Real-valued point :math:`x` at which to evaluate the log PDF.
+        :return: Log probability density evaluated under the Gumbel distribution.
+        """
         z = (value - self.loc) / self.scale
         return -(z + jnp.exp(-z)) - jnp.log(self.scale)
 
     @property
     def mean(self) -> ArrayLike:
+        r"""Mean of the Gumbel distribution:
+
+        .. math::
+            \mathbb{E}[X] = \mu + \beta \gamma
+
+        where :math:`\gamma \approx 0.5772\ldots` is the Euler-Mascheroni constant,
+        available at, :data:`~jax.numpy.euler_gamma`.
+        """
         return jnp.broadcast_to(
             self.loc + self.scale * jnp.euler_gamma, self.batch_shape
         )
 
     @property
     def variance(self) -> ArrayLike:
+        r"""Variance of the Gumbel distribution:
+
+        .. math::
+            \mathrm{Var}(X) = \frac{\pi^2}{6} \beta^2
+        """
         return jnp.broadcast_to(jnp.pi**2 / 6.0 * self.scale**2, self.batch_shape)
 
     def cdf(self, value: ArrayLike) -> ArrayLike:
+        r"""Cumulative Distribution Function (CDF) of the Gumbel distribution:
+
+        .. math::
+            F(x \mid \mu, \beta) = \exp\!\left(-\exp\!\left(-\frac{x - \mu}{\beta}\right)\right)
+
+        :param value: Real-valued point :math:`x` at which to evaluate the CDF.
+        :return: CDF values in :math:`[0, 1]`.
+        """
         return jnp.exp(-jnp.exp((self.loc - value) / self.scale))
 
     def icdf(self, q: ArrayLike) -> ArrayLike:
+        r"""Inverse CDF (quantile function) of the Gumbel distribution:
+
+        .. math::
+            F^{-1}(q \mid \mu, \beta) = \mu - \beta \ln(-\ln q),
+            \quad q \in (0, 1)
+
+        :param q: Quantile values in :math:`(0, 1)`.
+        :return: Real-valued quantiles of the Gumbel distribution at ``q``.
+        """
         return self.loc - self.scale * jnp.log(-jnp.log(q))
 
 
@@ -1415,6 +1485,22 @@ def variance(self) -> ArrayLike:
 
 
 class Laplace(Distribution):
+    r"""The Laplace (double-exponential) distribution, a continuous real-valued
+    distribution parameterized by location :math:`\mu` and scale :math:`b > 0`.
+    It is the distribution of the difference of two i.i.d. exponential variates
+    and has heavier tails than the Normal distribution.
+
+    The Probability Density Function (PDF) is:
+
+    .. math::
+        f(x \mid \mu, b) = \frac{1}{2 b}
+            \exp\!\left(-\frac{|x - \mu|}{b}\right),
+            \quad x \in \mathbb{R}
+
+    where :math:`\mu \in \mathbb{R}` is the location (:attr:`loc`) and
+    :math:`b > 0` is the scale (:attr:`scale`).
+    """
+
     arg_constraints = {"loc": constraints.real, "scale": constraints.positive}
     support = constraints.real
     reparametrized_params = ["loc", "scale"]
@@ -1426,13 +1512,26 @@ def __init__(
         *,
         validate_args: Optional[bool] = None,
     ) -> None:
+        r"""
+        :param loc: Location parameter :math:`\mu \in \mathbb{R}`. Defaults to ``0.0``.
+        :param scale: Scale parameter :math:`b > 0`. Defaults to ``1.0``.
+        :param validate_args: If True, enforce domain constraints during initialization.
+        """
         self.loc, self.scale = promote_shapes(loc, scale)
         batch_shape = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
         super(Laplace, self).__init__(
             batch_shape=batch_shape, validate_args=validate_args
         )
 
     def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+        r"""Draw samples via the location-scale transform
+        :math:`X = \mu + b Z`, where :math:`Z \sim \mathrm{Laplace}(0, 1)` is
+        drawn from :func:`~jax.random.laplace`.
+
+        :param key: A JAX PRNG key.
+        :param sample_shape: Sample dimensions to prepend to the batch shape.
+        :return: Real-valued samples from the Laplace distribution.
+        """
         assert is_prng_key(key)
         eps = random.laplace(
             key, shape=sample_shape + self.batch_shape + self.event_shape
@@ -1441,27 +1540,73 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
 
     @validate_sample
     def log_prob(self, value: ArrayLike) -> ArrayLike:
+        r"""Evaluate the log probability density function at ``value``:
+
+        .. math::
+            \ln f(x \mid \mu, b) = -\frac{|x - \mu|}{b} - \ln(2 b)
+
+        :param value: Real-valued point :math:`x` at which to evaluate the log PDF.
+        :return: Log probability density evaluated under the Laplace distribution.
+        """
         normalize_term = jnp.log(2 * self.scale)
         value_scaled = jnp.abs(value - self.loc) / self.scale
         return -value_scaled - normalize_term
 
     @property
     def mean(self) -> ArrayLike:
+        r"""Mean of the Laplace distribution:
+
+        .. math::
+            \mathbb{E}[X] = \mu
+        """
         return jnp.broadcast_to(self.loc, self.batch_shape)
 
     @property
     def variance(self) -> ArrayLike:
+        r"""Variance of the Laplace distribution:
+
+        .. math::
+            \mathrm{Var}(X) = 2 b^2
+        """
         return jnp.broadcast_to(2 * self.scale**2, self.batch_shape)
 
     def cdf(self, value: ArrayLike) -> ArrayLike:
+        r"""Cumulative Distribution Function (CDF) of the Laplace distribution.
+        Letting :math:`z = (x - \mu)/b`,
+
+        .. math::
+            F(x \mid \mu, b) = \frac{1}{2} - \frac{1}{2}\,
+                \operatorname{sgn}(z)\,\left(e^{-|z|} - 1\right)
+
+        The implementation uses :func:`~jax.numpy.expm1` for numerical
+        stability near :math:`z = 0`.
+
+        :param value: Real-valued point :math:`x` at which to evaluate the CDF.
+        :return: CDF values in :math:`[0, 1]`.
+        """
         scaled = (value - self.loc) / self.scale
         return 0.5 - 0.5 * jnp.sign(scaled) * jnp.expm1(-jnp.abs(scaled))
 
     def icdf(self, q: ArrayLike) -> ArrayLike:
+        r"""Inverse CDF (quantile function) of the Laplace distribution:
+
+        .. math::
+            F^{-1}(q \mid \mu, b) = \mu - b\,\mathrm{sgn}\left(q - \frac{1}{2}\right)\,
+                \ln\!\left(1 - 2 \left| q - \frac{1}{2} \right| \right),
+            \quad q \in [0, 1]
+
+        :param q: Quantile values in :math:`[0, 1]`.
+        :return: Real-valued quantiles of the Laplace distribution at ``q``.
+        """
         a = q - 0.5
         return self.loc - self.scale * jnp.sign(a) * jnp.log1p(-2 * jnp.abs(a))
 
     def entropy(self) -> ArrayLike:
+        r"""Differential entropy of the Laplace distribution:
+
+        .. math::
+            H(X) = \ln(2 b) + 1
+        """
         return jnp.log(2 * self.scale) + 1
 
 
@@ -1782,6 +1927,25 @@ def entropy(self) -> ArrayLike:
 
 
 class Logistic(Distribution):
+    r"""The Logistic distribution, a continuous real-valued distribution
+    parameterized by location :math:`\mu` and scale :math:`s > 0`. Its CDF is
+    the standard logistic (sigmoid) function shifted and scaled to :math:`\mu`,
+    :math:`s`, which makes it the natural latent distribution underlying
+    logistic regression.
+
+    The Probability Density Function (PDF) is:
+
+    .. math::
+        f(x \mid \mu, s) = \frac{
+            \exp\!\left(-\displaystyle\frac{x - \mu}{s}\right)
+        }{
+            s \left(1 + \exp\!\left(-\displaystyle\frac{x - \mu}{s}\right)\right)^{2}
+        }, \quad x \in \mathbb{R}
+
+    where :math:`\mu \in \mathbb{R}` is the location (:attr:`loc`) and
+    :math:`s > 0` is the scale (:attr:`scale`).
+    """
+
     arg_constraints = {"loc": constraints.real, "scale": constraints.positive}
     support = constraints.real
     reparametrized_params = ["loc", "scale"]
@@ -1793,11 +1957,24 @@ def __init__(
         *,
         validate_args: Optional[bool] = None,
     ) -> None:
+        r"""
+        :param loc: Location parameter :math:`\mu \in \mathbb{R}`. Defaults to ``0.0``.
+        :param scale: Scale parameter :math:`s > 0`. Defaults to ``1.0``.
+        :param validate_args: If True, enforce domain constraints during initialization.
+        """
         self.loc, self.scale = promote_shapes(loc, scale)
         batch_shape = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
         super(Logistic, self).__init__(batch_shape, validate_args=validate_args)
 
     def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
+        r"""Draw samples via the location-scale transform
+        :math:`X = \mu + s Z`, where :math:`Z \sim \mathrm{Logistic}(0, 1)` is
+        drawn from :func:`~jax.random.logistic`.
+
+        :param key: A JAX PRNG key.
+        :param sample_shape: Sample dimensions to prepend to the batch shape.
+        :return: Real-valued samples from the Logistic distribution.
+        """
         assert is_prng_key(key)
         z = random.logistic(
             key, shape=sample_shape + self.batch_shape + self.event_shape
@@ -1806,27 +1983,79 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
 
     @validate_sample
     def log_prob(self, value: ArrayLike) -> ArrayLike:
+        r"""Evaluate the log probability density function at ``value``.
+
+        Letting :math:`u = (\mu - x)/s`, the log PDF is
+
+        .. math::
+            \ln f(x \mid \mu, s) = u - \ln s - 2 \ln(1 + e^{u})
+
+        The implementation uses :func:`~jax.nn.softplus` for
+        :math:`\ln(1 + e^{u})`, which is numerically stable for both large
+        positive and large negative values of :math:`u`.
+
+        :param value: Real-valued point :math:`x` at which to evaluate the log PDF.
+        :return: Log probability density evaluated under the Logistic distribution.
+        """
         log_exponent = (self.loc - value) / self.scale
         log_denominator = jnp.log(self.scale) + 2 * nn.softplus(log_exponent)
         return log_exponent - log_denominator
 
     @property
     def mean(self) -> ArrayLike:
+        r"""Mean of the Logistic distribution:
+
+        .. math::
+            \mathbb{E}[X] = \mu
+        """
         return jnp.broadcast_to(self.loc, self.batch_shape)
 
     @property
     def variance(self) -> ArrayLike:
+        r"""Variance of the Logistic distribution:
+
+        .. math::
+            \mathrm{Var}(X) = \frac{\pi^2 s^2}{3}
+        """
         var = (self.scale**2) * (jnp.pi**2) / 3
         return jnp.broadcast_to(var, self.batch_shape)
 
     def cdf(self, value: ArrayLike) -> ArrayLike:
+        r"""Cumulative Distribution Function (CDF) of the Logistic distribution.
+        Letting :math:`z = (x - \mu)/s`,
+
+        .. math::
+            F(x \mid \mu, s) = \sigma(z) = \frac{1}{1 + e^{-z}}
+
+        where :math:`\sigma` is the logistic sigmoid, computed via
+        :func:`~jax.scipy.special.expit`.
+
+        :param value: Real-valued point :math:`x` at which to evaluate the CDF.
+        :return: CDF values in :math:`[0, 1]`.
+        """
         scaled = (value - self.loc) / self.scale
         return expit(scaled)
 
     def icdf(self, q: ArrayLike) -> ArrayLike:
+        r"""Inverse CDF (quantile function) of the Logistic distribution:
+
+        .. math::
+            F^{-1}(q \mid \mu, s) = \mu + s\,\operatorname{logit}(q),
+            \quad q \in [0, 1]
+
+        where :math:`\operatorname{logit}(q) = \ln(q / (1 - q))`.
+
+        :param q: Quantile values in :math:`[0, 1]`.
+        :return: Real-valued quantiles of the Logistic distribution at ``q``.
+        """
         return self.loc + self.scale * logit(q)
 
     def entropy(self) -> ArrayLike:
+        r"""Differential entropy of the Logistic distribution:
+
+        .. math::
+            H(X) = \ln(s) + 2
+        """
         return jnp.broadcast_to(jnp.log(self.scale) + 2, self.batch_shape)