Return floats instead of integers in ShiftedBetaGeometricRV.rng_fn

pymc_marketing/clv/distributions.py

@@ -897,7 +897,7 @@ class ShiftedBetaGeometricRV(RandomVariable):
     name = "sbg"
     signature = "(),()->()"
 
-    dtype = "int64"
+    dtype = "floatX"
     _print_name = ("ShiftedBetaGeometric", "\\operatorname{ShiftedBetaGeometric}")
 
     @classmethod
@@ -912,9 +912,9 @@ def rng_fn(cls, rng, alpha, beta, size):
 
         # prevent log(0) by clipping small p samples
         p = np.clip(p_samples, 1e-100, 1)
-        # TODO: Consider returning np.float64 types instead of np.int64
-        #       See relevant PR comment: https://github.com/pymc-labs/pymc-marketing/pull/2010#discussion_r2444116986
-        return rng.geometric(p, size=size)
+        # Convert to float to match dtype="floatX" specification
+        # Using numpy's default float type for consistency with other distributions
+        return rng.geometric(p, size=size).astype(np.float64)
 
 
 sbg = ShiftedBetaGeometricRV()

tests/clv/test_distributions.py

@@ -824,6 +824,16 @@ def test_random_edge_cases(self):
                     assert np.mean(draws) >= 1
                     assert np.var(draws) >= 0
 
+        def test_random_returns_float_type(self):
+            """Test that rng_fn returns float type, not integer type"""
+            dist = self.pymc_dist.dist(alpha=1.0, beta=1.0, size=10)
+            draws = dist.eval()
+
+            # Verify that the dtype is float, not int
+            assert np.issubdtype(draws.dtype, np.floating), (
+                f"Expected float dtype, but got {draws.dtype}"
+            )
+
     def test_logp(self):
         alpha = pt.scalar("alpha")
         beta = pt.scalar("beta")