fix(gh-2181): ensure consistent behavior of Poisson.log_prob by casting rate to floating-point type (#2182)

* fix(gh-2181): ensure consistent behavior of `Poisson.log_prob` by casting rate to floating-point type

* test: unit test that covers the fix

* fix: remove float values from test cases

Author: Qazalbash

numpyro/distributions/discrete.py

@@ -783,32 +783,33 @@ def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLik
 
     @validate_sample
     def log_prob(self, value: ArrayLike) -> ArrayLike:
-        if self._validate_args:
-            self._validate_sample(value)
+        # Using an integer vs. floating-point `rate` leads to differing results.
+        # To ensure consistent behavior, `rate` is explicitly cast to a floating-point type.
+        # See: https://github.com/pyro-ppl/numpyro/issues/2181
+        ftype = jnp.result_type(float)
+        rate = jnp.astype(self.rate, ftype)
+
         if (
             self.is_sparse
             and not isinstance(value, jax.core.Tracer)
             and jnp.size(value) > 1
         ):
             shape = lax.broadcast_shapes(self.batch_shape, jnp.shape(value))
-            rate = jnp.broadcast_to(self.rate, shape).reshape(-1)
+            rate = jnp.broadcast_to(rate, shape).reshape(-1)
             nonzero = np.broadcast_to(jax.device_get(value) > 0, shape).reshape(-1)
             value = jnp.broadcast_to(value, shape).reshape(-1)
             sparse_value = value[nonzero]
             sparse_rate = rate[nonzero]
             return (
-                jnp.asarray(-rate, dtype=jnp.result_type(float))
+                jnp.asarray(-rate, dtype=ftype)
                 .at[nonzero]
                 .add(
                     jnp.log(sparse_rate) * sparse_value - gammaln(sparse_value + 1),
                 )
                 .reshape(shape)
             )
-        return (
-            xlogy(jnp.astype(value, jnp.result_type(self.rate)), self.rate)
-            - gammaln(value + 1)
-            - self.rate
-        )
+        _value = jnp.astype(value, ftype)
+        return xlogy(_value, rate) - gammaln(_value + 1.0) - rate
 
     @property
     def mean(self) -> ArrayLike:

test/test_distributions.py

@@ -4810,3 +4810,38 @@ def test_uniform_log_prob_outside_support():
         match="Out-of-support values provided to log prob method. The value argument should be within the support.",
     ):
         d.log_prob(-0.5)
+
+
+@pytest.mark.parametrize("rate", [0, 1, 2, 5, 10, 1e-6, 1e2])
+@pytest.mark.parametrize("value", [0, 1, 2, 5, 10])
+def test_poisson_dtype_consistency(rate, value):
+    """
+    Ensure that ``Poisson.log_prob`` is invariant to dtype differences in both
+    the rate parameter and the observed value.
+
+    This test checks that using integer vs. floating-point representations for:
+      - the rate (e.g., ``2`` vs ``2.0``), and
+      - the observed value (e.g., ``2`` vs ``2.0``),
+
+    yields identical log-probabilities across all combinations. This includes
+    edge cases such as zero rate, very small rates, large rates, and values
+    near integer boundaries.
+
+    The test guards against dtype-dependent behavior, where numerically
+    equivalent inputs produce inconsistent results due to type casting or
+    implementation details.
+
+    See: https://github.com/pyro-ppl/numpyro/issues/2181
+    """
+    rates = [rate, float(rate)]
+
+    results = []
+    for r in rates:
+        d = dist.Poisson(r, validate_args=True)
+        results.append(d.log_prob(value))
+
+    ref = results[0]
+    for res in results[1:]:
+        assert jnp.allclose(res, ref), (
+            f"Inconsistent results for rate={rate}, value={value}: {results}"
+        )