Cast rounded input replacement to original discrete dtype in sample_smc

Author: ricardoV94

RELEASE-NOTES.md

@@ -3,7 +3,7 @@
 ## PyMC 4.0.1 (vNext)
 + Fixed an incorrect entry in `pm.Metropolis.stats_dtypes` (see #5582).
 + Added a check in `Empirical` approximation which does not yet support `InferenceData` inputs (see #5874, #5884).
-+ ...
++ Fixed bug when sampling discrete variables with SMC (see #5887).
 
 ## PyMC 4.0.0 (2022-06-03)
 

pymc/smc/smc.py

@@ -567,7 +567,7 @@ def _logp_forw(point, out_vars, in_vars, shared):
             if in_var.dtype in discrete_types:
                 float_var = at.TensorType("floatX", in_var.broadcastable)(in_var.name)
                 new_in_vars.append(float_var)
-                replace_int_input[in_var] = at.round(float_var)
+                replace_int_input[in_var] = at.round(float_var).astype(in_var.dtype)
             else:
                 new_in_vars.append(in_var)
 

pymc/tests/test_smc.py

@@ -110,7 +110,7 @@ def test_discrete_rounding_proposal(self):
         assert np.isclose(smc.prior_logp_func(floatX(np.array([0.51]))), np.log(0.7))
         assert smc.prior_logp_func(floatX(np.array([1.51]))) == -np.inf
 
-    def test_unobserved_discrete(self):
+    def test_unobserved_bernoulli(self):
         n = 10
         rng = self.get_random_state()
         z_true = np.zeros(n, dtype=int)
@@ -126,6 +126,15 @@ def test_unobserved_discrete(self):
 
         assert np.all(np.median(trace["z"], axis=0) == z_true)
 
+    def test_unobserved_categorical(self):
+        with pm.Model() as m:
+            mu = pm.Categorical("mu", p=[0.1, 0.3, 0.6], size=2)
+            pm.Normal("like", mu=mu, sigma=0.1, observed=[1, 2])
+
+            trace = pm.sample_smc(chains=1, return_inferencedata=False)
+
+        assert np.all(np.median(trace["mu"], axis=0) == [1, 2])
+
     def test_marginal_likelihood(self):
         """
         Verifies that the log marginal likelihood function