Add Poisson and NegativeBinomial to pymc.dims (#8305)

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

Author: williambdean

pymc/dims/distributions/scalar.py

@@ -36,6 +36,7 @@
     halfflat,
     truncated_normal,
 )
+from pymc.distributions.discrete import NegativeBinomial as RegularNegativeBinomial
 from pymc.util import UNSET
 
 
@@ -295,3 +296,22 @@ def xrv_op(self, alpha, beta, core_dims=None, extra_dims=None, rng=None, **kwarg
         core_rv = WeibullBetaRV.rv_op(alpha=alpha.values, beta=beta.values).owner.op
         xop = ptxr.as_xrv(core_rv)
         return xop(alpha, beta, core_dims=core_dims, extra_dims=extra_dims, rng=rng, **kwargs)
+
+
+@copy_docstring(regular_dists.Poisson)
+class Poisson(DimDistribution):
+    xrv_op = ptxr.poisson
+
+    @classmethod
+    def dist(cls, mu, **kwargs):
+        return super().dist([mu], **kwargs)
+
+
+@copy_docstring(regular_dists.NegativeBinomial)
+class NegativeBinomial(DimDistribution):
+    xrv_op = ptxr.nbinom
+
+    @classmethod
+    def dist(cls, mu=None, alpha=None, *, p=None, n=None, **kwargs):
+        n, p = RegularNegativeBinomial.get_n_p(mu=mu, alpha=alpha, p=p, n=n)
+        return super().dist([n, p], **kwargs)

tests/dims/distributions/test_scalar.py

@@ -31,7 +31,9 @@
     InverseGamma,
     Laplace,
     LogNormal,
+    NegativeBinomial,
     Normal,
+    Poisson,
     StudentT,
     TruncatedNormal,
     Uniform,
@@ -311,3 +313,27 @@ def test_weibull():
 
     assert_equivalent_random_graph(model, reference_model)
     assert_equivalent_logp_graph(model, reference_model)
+
+
+def test_poisson():
+    coords = {"a": range(3)}
+    with Model(coords=coords) as model:
+        Poisson("x", mu=2.0, dims="a")
+
+    with Model(coords=coords) as reference_model:
+        regular_distributions.Poisson("x", mu=2.0, dims="a")
+
+    assert_equivalent_random_graph(model, reference_model)
+    assert_equivalent_logp_graph(model, reference_model)
+
+
+def test_negative_binomial():
+    coords = {"a": range(3)}
+    with Model(coords=coords) as model:
+        NegativeBinomial("x", mu=5.0, alpha=2.0, dims="a")
+
+    with Model(coords=coords) as reference_model:
+        regular_distributions.NegativeBinomial("x", mu=5.0, alpha=2.0, dims="a")
+
+    assert_equivalent_random_graph(model, reference_model)
+    assert_equivalent_logp_graph(model, reference_model)