Merge upstream/main to incorporate latest fixes

Author: sotagg

pymc/dims/distributions/transforms.py

@@ -54,6 +54,34 @@ def log_jac_det(self, value, *inputs):
 log_odds_transform = LogOddsTransform()
 
 
+class SimplexTransform(DimTransform):
+    name = "simplex"
+
+    def __init__(self, dim: str):
+        self.core_dim = dim
+
+    def forward(self, value, *inputs):
+        log_value = ptx.math.log(value)
+        N = value.sizes[self.core_dim].astype(value.dtype)
+        shift = log_value.sum(self.core_dim) / N
+        return log_value.isel({self.core_dim: slice(None, -1)}) - shift
+
+    def backward(self, value, *inputs):
+        value = ptx.concat([value, -value.sum(self.core_dim)], dim=self.core_dim)
+        exp_value_max = ptx.math.exp(value - value.max(self.core_dim))
+        return exp_value_max / exp_value_max.sum(self.core_dim)
+
+    def log_jac_det(self, value, *inputs):
+        N = value.sizes[self.core_dim] + 1
+        N = N.astype(value.dtype)
+        sum_value = value.sum(self.core_dim)
+        value_sum_expanded = value + sum_value
+        value_sum_expanded = ptx.concat([value_sum_expanded, 0], dim=self.core_dim)
+        logsumexp_value_expanded = ptx.math.logsumexp(value_sum_expanded, dim=self.core_dim)
+        res = ptx.math.log(N) + (N * sum_value) - (N * logsumexp_value_expanded)
+        return res
+
+
 class ZeroSumTransform(DimTransform):
     name = "zerosum"
 

pymc/dims/distributions/vector.py

@@ -19,7 +19,7 @@
 from pytensor.xtensor import random as pxr
 
 from pymc.dims.distributions.core import VectorDimDistribution
-from pymc.dims.distributions.transforms import ZeroSumTransform
+from pymc.dims.distributions.transforms import SimplexTransform, ZeroSumTransform
 from pymc.distributions.multivariate import ZeroSumNormalRV
 from pymc.util import UNSET
 
@@ -63,6 +63,61 @@ def dist(cls, p=None, *, logit_p=None, core_dims=None, **kwargs):
         return super().dist([p], core_dims=core_dims, **kwargs)
 
 
+class Dirichlet(VectorDimDistribution):
+    """Dirichlet distribution.
+
+    Parameters
+    ----------
+    a : xtensor_like, optional
+        Probabilities of each category. Must sum to 1 along the core dimension.
+    core_dims : str
+        The core dimension of the distribution, which represents the categories.
+        The dimension must be present in `p` or `logit_p`.
+    **kwargs
+        Other keyword arguments used to define the distribution.
+
+    Returns
+    -------
+    XTensorVariable
+        An xtensor variable representing the categorical distribution.
+        The output does not contain the core dimension, as it is absorbed into the distribution.
+
+
+    """
+
+    xrv_op = ptxr.dirichlet
+
+    @classmethod
+    def __new__(
+        cls, *args, core_dims=None, dims=None, default_transform=UNSET, observed=None, **kwargs
+    ):
+        if core_dims is not None:
+            if isinstance(core_dims, tuple | list):
+                [core_dims] = core_dims
+
+            # Create default_transform
+            if observed is None and default_transform is UNSET:
+                default_transform = SimplexTransform(dim=core_dims)
+
+        # If the user didn't specify dims, take it from core_dims
+        # We need them to be forwarded to dist in the `dim_lenghts` argument
+        # if dims is None and core_dims is not None:
+        #    dims = (..., *core_dims)
+
+        return super().__new__(
+            *args,
+            core_dims=core_dims,
+            dims=dims,
+            default_transform=default_transform,
+            observed=observed,
+            **kwargs,
+        )
+
+    @classmethod
+    def dist(cls, a, *, core_dims=None, **kwargs):
+        return super().dist([a], core_dims=core_dims, **kwargs)
+
+
 class MvNormal(VectorDimDistribution):
     """Multivariate Normal distribution.
 

pymc/distributions/continuous.py

@@ -2575,6 +2575,11 @@ class ChiSquared:
     ----------
     nu : tensor_like of float
         Degrees of freedom (nu > 0).
+
+    Notes
+    -----
+    This is implemented as a special case of the Gamma distribution.
+    :math:`\chi^2(\nu) = \text{Gamma}(\alpha=\nu/2, \beta=1/2)`
     """
 
     def __new__(cls, name, nu, **kwargs):
@@ -3601,7 +3606,7 @@ def icdf(value, mu, s):
 class LogitNormalRV(SymbolicRandomVariable):
     name = "logit_normal"
     extended_signature = "[rng],[size],(),()->[rng],()"
-    _print_name = ("logitNormal", "\\operatorname{logitNormal}")
+    _print_name = ("LogitNormal", "\\operatorname{LogitNormal}")
 
     @classmethod
     def rv_op(cls, mu, sigma, *, size=None, rng=None):

pymc/distributions/discrete.py

@@ -393,7 +393,7 @@ def logcdf(value, p):
 class DiscreteWeibullRV(SymbolicRandomVariable):
     name = "discrete_weibull"
     extended_signature = "[rng],[size],(),()->[rng],()"
-    _print_name = ("dWeibull", "\\operatorname{dWeibull}")
+    _print_name = ("DiscreteWeibull", "\\operatorname{DiscreteWeibull}")
 
     @classmethod
     def rv_op(cls, q, beta, *, size=None, rng=None):

tests/dims/distributions/test_vector.py

@@ -19,7 +19,7 @@
 import pymc.distributions as regular_distributions
 
 from pymc import Model
-from pymc.dims import Categorical, MvNormal, ZeroSumNormal
+from pymc.dims import Categorical, Dirichlet, MvNormal, ZeroSumNormal
 from tests.dims.utils import assert_equivalent_logp_graph, assert_equivalent_random_graph
 
 
@@ -40,6 +40,27 @@ def test_categorical():
     assert_equivalent_logp_graph(model, reference_model)
 
 
+def test_dirichlet():
+    coords = {"a": range(3), "b": range(2)}
+    alpha = pt.as_tensor([1, 2, 3])
+
+    alpha_xr = as_xtensor(alpha, dims=("b",))
+
+    with Model(coords=coords) as model:
+        Dirichlet("x", a=alpha_xr, core_dims="b", dims=("a", "b"))
+
+    with Model(coords=coords) as reference_model:
+        regular_distributions.Dirichlet("x", a=alpha, dims=("a", "b"))
+
+    assert_equivalent_random_graph(model, reference_model)
+
+    # logp graphs end up different, but they mean the same thing
+    np.testing.assert_allclose(
+        model.compile_logp()(model.initial_point()),
+        reference_model.compile_logp()(reference_model.initial_point()),
+    )
+
+
 def test_mvnormal():
     coords = {"a": range(3), "b": range(2)}
     mu = pt.as_tensor([1, 2])

tests/distributions/test_censored.py

@@ -137,15 +137,13 @@ def test_censored_logcdf_continuous(self):
 
         # No censoring
         censored_norm = pm.Censored.dist(norm, lower=None, upper=None)
-        with pytest.warns(RuntimeWarning, match=match_str):
-            censored_eval = logcdf(censored_norm, eval_points).eval()
+        censored_eval = logcdf(censored_norm, eval_points).eval()
         np.testing.assert_allclose(censored_eval, expected_logcdf_uncensored)
 
         # Left censoring
         censored_norm = pm.Censored.dist(norm, lower=-1, upper=None)
         expected_left = np.where(eval_points < -1, -np.inf, expected_logcdf_uncensored)
-        with pytest.warns(RuntimeWarning, match=match_str):
-            censored_eval = logcdf(censored_norm, eval_points).eval()
+        censored_eval = logcdf(censored_norm, eval_points).eval()
         np.testing.assert_allclose(
             censored_eval,
             expected_left,
@@ -155,8 +153,7 @@ def test_censored_logcdf_continuous(self):
         # Right censoring
         censored_norm = pm.Censored.dist(norm, lower=None, upper=1)
         expected_right = np.where(eval_points >= 1, 0.0, expected_logcdf_uncensored)
-        with pytest.warns(RuntimeWarning, match=match_str):
-            censored_eval = logcdf(censored_norm, eval_points).eval()
+        censored_eval = logcdf(censored_norm, eval_points).eval()
         np.testing.assert_allclose(
             censored_eval,
             expected_right,
@@ -167,8 +164,7 @@ def test_censored_logcdf_continuous(self):
         censored_norm = pm.Censored.dist(norm, lower=-1, upper=1)
         expected_interval = np.where(eval_points < -1, -np.inf, expected_logcdf_uncensored)
         expected_interval = np.where(eval_points >= 1, 0.0, expected_interval)
-        with pytest.warns(RuntimeWarning, match=match_str):
-            censored_eval = logcdf(censored_norm, eval_points).eval()
+        censored_eval = logcdf(censored_norm, eval_points).eval()
         np.testing.assert_allclose(
             censored_eval,
             expected_interval,