Allow Constant to be integer or float

Author: ricardoV94

pymc/distributions/discrete.py

@@ -46,6 +46,7 @@
 from pymc.distributions.logprob import logcdf, logp
 from pymc.distributions.shape_utils import rv_size_is_none
 from pymc.math import sigmoid
+from pymc.vartypes import continuous_types
 
 __all__ = [
     "Binomial",
@@ -1315,9 +1316,12 @@ class ConstantRV(RandomVariable):
     name = "constant"
     ndim_supp = 0
     ndims_params = [0]
-    dtype = "floatX"  # Should be treated as a discrete variable!
     _print_name = ("Constant", "\\operatorname{Constant}")
 
+    def make_node(self, rng, size, dtype, c):
+        c = at.as_tensor_variable(c)
+        return super().make_node(rng, size, c.dtype, c)
+
     @classmethod
     def rng_fn(cls, rng, c, size=None):
         if size is None:
@@ -1334,15 +1338,19 @@ class Constant(Discrete):
 
     Parameters
     ----------
-    value: float or int
-        Constant parameter.
+    c: float or int
+        Constant parameter. The dtype of `c` determines the dtype of the distribution.
+        This can affect which sampler is assigned to Constant variables, or variables
+        that use Constant, such as Mixtures.
     """
 
     rv_op = constant
 
     @classmethod
     def dist(cls, c, *args, **kwargs):
-        c = at.as_tensor_variable(floatX(c))
+        c = at.as_tensor_variable(c)
+        if c.dtype in continuous_types:
+            c = floatX(c)
         return super().dist([c], **kwargs)
 
     def get_moment(rv, size, c):

pymc/tests/test_distributions.py

@@ -130,7 +130,7 @@ def polyagamma_cdf(*args, **kwargs):
 from pymc.math import kronecker
 from pymc.model import Deterministic, Model, Point, Potential
 from pymc.tests.helpers import select_by_precision
-from pymc.vartypes import continuous_types
+from pymc.vartypes import continuous_types, discrete_types
 
 
 def get_lkj_cases():
@@ -934,7 +934,9 @@ def check_selfconsistency_discrete_logcdf(
         Check that logcdf of discrete distributions matches sum of logps up to value
         """
         # This test only works for scalar random variables
-        assert distribution.rv_op.ndim_supp == 0
+        rv_op = getattr(distribution, "rv_op", None)
+        if rv_op:
+            assert rv_op.ndim_supp == 0
 
         domains = paramdomains.copy()
         domains["value"] = domain
@@ -3416,3 +3418,17 @@ def test_sd_dist_automatically_resized(self, sd_dist):
         assert resized_sd_dist.eval().shape == (10, 3)
         # LKJCov has support shape `(n * (n+1)) // 2`
         assert x.eval().shape == (10, 6)
+
+
+@pytest.mark.parametrize(
+    "dist, non_psi_args",
+    [
+        (pm.ZeroInflatedPoisson.dist, (2,)),
+        (pm.ZeroInflatedBinomial.dist, (2, 0.5)),
+        (pm.ZeroInflatedNegativeBinomial.dist, (2, 2)),
+    ],
+)
+def test_zero_inflated_dists_dtype_and_broadcast(dist, non_psi_args):
+    x = dist([0.5, 0.5, 0.5], *non_psi_args)
+    assert x.dtype in discrete_types
+    assert x.eval().shape == (3,)

pymc/tests/test_distributions_moments.py

@@ -626,7 +626,7 @@ def test_constant_moment(c, size, expected):
 @pytest.mark.parametrize(
     "psi, theta, size, expected",
     [
-        (0.9, 3.0, None, 2),
+        (0.9, 3.0, None, 3),
         (0.8, 2.9, 5, np.full(5, 2)),
         (0.2, np.arange(1, 5) * 5, None, np.arange(1, 5)),
         (0.2, np.arange(1, 5) * 5, (2, 4), np.full((2, 4), np.arange(1, 5))),
@@ -1335,7 +1335,13 @@ def test_multinomial_moment(p, n, size, expected):
     [
         (0.2, 10, 3, None, 2),
         (0.2, 10, 4, 5, np.full(5, 2)),
-        (0.4, np.arange(1, 5), np.arange(2, 6), None, np.array([0, 0, 1, 1])),
+        (
+            0.4,
+            np.arange(1, 5),
+            np.arange(2, 6),
+            None,
+            np.array([0, 1, 1, 2] if aesara.config.floatX == "float64" else [0, 0, 1, 1]),
+        ),
         (
             np.linspace(0.2, 0.6, 3),
             np.arange(1, 10, 4),

pymc/tests/test_distributions_random.py

@@ -1571,112 +1571,14 @@ def constant_rng_fn(self, size, c):
         "check_pymc_params_match_rv_op",
         "check_pymc_draws_match_reference",
         "check_rv_size",
+        "check_dtype",
     ]
 
-
-class TestZeroInflatedPoisson(BaseTestDistributionRandom):
-    def zero_inflated_poisson_rng_fn(self, size, psi, theta, poisson_rng_fct, random_rng_fct):
-        return poisson_rng_fct(theta, size=size) * (random_rng_fct(size=size) < psi)
-
-    def seeded_zero_inflated_poisson_rng_fn(self):
-        poisson_rng_fct = functools.partial(
-            getattr(np.random.RandomState, "poisson"), self.get_random_state()
-        )
-
-        random_rng_fct = functools.partial(
-            getattr(np.random.RandomState, "random"), self.get_random_state()
-        )
-
-        return functools.partial(
-            self.zero_inflated_poisson_rng_fn,
-            poisson_rng_fct=poisson_rng_fct,
-            random_rng_fct=random_rng_fct,
-        )
-
-    pymc_dist = pm.ZeroInflatedPoisson
-    pymc_dist_params = {"psi": 0.9, "theta": 4.0}
-    expected_rv_op_params = {"psi": 0.9, "theta": 4.0}
-    reference_dist_params = {"psi": 0.9, "theta": 4.0}
-    reference_dist = seeded_zero_inflated_poisson_rng_fn
-    checks_to_run = [
-        "check_pymc_params_match_rv_op",
-        "check_pymc_draws_match_reference",
-        "check_rv_size",
-    ]
-
-
-class TestZeroInflatedBinomial(BaseTestDistributionRandom):
-    def zero_inflated_binomial_rng_fn(self, size, psi, n, p, binomial_rng_fct, random_rng_fct):
-        return binomial_rng_fct(n, p, size=size) * (random_rng_fct(size=size) < psi)
-
-    def seeded_zero_inflated_binomial_rng_fn(self):
-        binomial_rng_fct = functools.partial(
-            getattr(np.random.RandomState, "binomial"), self.get_random_state()
-        )
-
-        random_rng_fct = functools.partial(
-            getattr(np.random.RandomState, "random"), self.get_random_state()
-        )
-
-        return functools.partial(
-            self.zero_inflated_binomial_rng_fn,
-            binomial_rng_fct=binomial_rng_fct,
-            random_rng_fct=random_rng_fct,
-        )
-
-    pymc_dist = pm.ZeroInflatedBinomial
-    pymc_dist_params = {"psi": 0.9, "n": 12, "p": 0.7}
-    expected_rv_op_params = {"psi": 0.9, "n": 12, "p": 0.7}
-    reference_dist_params = {"psi": 0.9, "n": 12, "p": 0.7}
-    reference_dist = seeded_zero_inflated_binomial_rng_fn
-    checks_to_run = [
-        "check_pymc_params_match_rv_op",
-        "check_pymc_draws_match_reference",
-        "check_rv_size",
-    ]
-
-
-class TestZeroInflatedNegativeBinomialMuSigma(BaseTestDistributionRandom):
-    def zero_inflated_negbinomial_rng_fn(
-        self, size, psi, n, p, negbinomial_rng_fct, random_rng_fct
-    ):
-        return negbinomial_rng_fct(n, p, size=size) * (random_rng_fct(size=size) < psi)
-
-    def seeded_zero_inflated_negbinomial_rng_fn(self):
-        negbinomial_rng_fct = functools.partial(
-            getattr(np.random.RandomState, "negative_binomial"), self.get_random_state()
-        )
-
-        random_rng_fct = functools.partial(
-            getattr(np.random.RandomState, "random"), self.get_random_state()
-        )
-
-        return functools.partial(
-            self.zero_inflated_negbinomial_rng_fn,
-            negbinomial_rng_fct=negbinomial_rng_fct,
-            random_rng_fct=random_rng_fct,
-        )
-
-    n, p = pm.NegativeBinomial.get_n_p(mu=3, alpha=5)
-
-    pymc_dist = pm.ZeroInflatedNegativeBinomial
-    pymc_dist_params = {"psi": 0.9, "mu": 3, "alpha": 5}
-    expected_rv_op_params = {"psi": 0.9, "n": n, "p": p}
-    reference_dist_params = {"psi": 0.9, "n": n, "p": p}
-    reference_dist = seeded_zero_inflated_negbinomial_rng_fn
-    checks_to_run = [
-        "check_pymc_params_match_rv_op",
-        "check_pymc_draws_match_reference",
-        "check_rv_size",
-    ]
-
-
-class TestZeroInflatedNegativeBinomial(BaseTestDistributionRandom):
-    pymc_dist = pm.ZeroInflatedNegativeBinomial
-    pymc_dist_params = {"psi": 0.9, "n": 12, "p": 0.7}
-    expected_rv_op_params = {"psi": 0.9, "n": 12, "p": 0.7}
-    reference_dist_params = {"psi": 0.9, "n": 12, "p": 0.7}
-    checks_to_run = ["check_pymc_params_match_rv_op"]
+    def check_dtype(self):
+        assert pm.Constant.dist(2**4).dtype == "int8"
+        assert pm.Constant.dist(2**16).dtype == "int32"
+        assert pm.Constant.dist(2**32).dtype == "int64"
+        assert pm.Constant.dist(2.0).dtype == aesara.config.floatX
 
 
 class TestOrderedLogistic(BaseTestDistributionRandom):