Cast rate to float when constructing Poisson to ensure consistent log prob behavior

[dylanhmorris]

Very excited about #2176. Wanted to propose a tweak to address a potential edge case. As written, this will cast real value arrays to integer arrays when the Poisson distribution instance was initialized with an integer rate.

numpyro/numpyro/distributions/discrete.py

Lines 807 to 811 in e708f34

	return (
	xlogy(jnp.astype(value, jnp.result_type(self.rate)), self.rate)
	- gammaln(value + 1)
	- self.rate
	)

That in turns means that (for example) Poisson(2).log_prob(value) and Poisson(2.0).log_prob(value) can behave differently for the same value.

I think we should avoid this by casting all rates to floats at construction. Rationale below. Let me know what you think, @fehiepsi, @Qazalbash, and @tillahoffmann.

Rationale

jnp.astype (mostly) casts floats to ints as the floor, e.g.

import jax.numpy as jnp
jnp.astype(1.99999, "int32")
# Array(1, dtype=int32)

Even though the Poisson is supported only on the integers, the log_prob method evaluates the PMF expression for non-integer values (as opposed to returning -inf). For example:

import numpyro.distributions as dist
my_poisson = dist.Poisson(2.5)
my_poisson.log_prob(5.2)
# Array(-2.8675885, dtype=float32)

Notably, my_poisson.log_prob(0.9999) and my_poisson.log_prob(1) give very similar results.

Now suppose we construct a Poisson with an integer rate, e.g. Poisson(2). Depending on whether that rate was actually an int (2) or was a float (2., 2.0, etc.) evaluating log_prob can give very different results for the same near-integer value.

int_rate = dist.Poisson(2)
float_rate = dist.Poisson(2.)
int_rate.log_prob(2)
# Array(-1.306853, dtype=float32)
float_rate.log_prob(2)
# Array(-1.306853, dtype=float32)
int_rate.log_prob(1.999999)
# Array(-2.0000002, dtype=float32)
float_rate.log_prob(1.999999)
# Array(-1.3068538, dtype=float32)

[Qazalbash]

Hi @dylanhmorris, I have proposed a fix. Please see #2182.

[Qazalbash]

Interesting behaviour when validation is on!

In [1]: import numpyro.distributions as dist

In [2]: int_rate = dist.Poisson(2, validate_args=True)

In [3]: float_rate = dist.Poisson(2., validate_args=True)

In [4]: int_rate.log_prob(2)
Out[4]: Array(-1.306853, dtype=float32)

In [5]: float_rate.log_prob(2)
Out[5]: Array(-1.306853, dtype=float32)

In [6]: int_rate.log_prob(1.999999)
<ipython-input-6-37287a8a65a9>:1: UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
  int_rate.log_prob(1.999999)
Out[6]: Array(-inf, dtype=float32)

In [7]: float_rate.log_prob(1.999999)
<ipython-input-7-119894d57215>:1: UserWarning: Out-of-support values provided to log prob method. The value argument should be within the support.
  float_rate.log_prob(1.999999)
Out[7]: Array(-inf, dtype=float32)

@dylanhmorris, we can avoid the issue by enabling the validation 😄.

[dylanhmorris]

Good point! I still think it's worth changing the behavior to ensure things work more predictably with validate_args=False, but I defer to others.

[tillahoffmann]

This is an interesting one. I agree that we should probably cast to the appropriate dtype for the constraint of the parameter because the behavior can be surprising. E.g., passing an int covariance matrix to a multivariate normal leaves the passed parameter as int, but the Cholesky decomposition has floating dtype.

In [1]: from numpyro import distributions as dists
        from jax import numpy as jnp

In [2]: def get_args(dist):
   ...:     return {key: getattr(dist, key, None) for key in dist.arg_constraints}
   ...: 

In [10]: get_args(dists.MultivariateNormal(0, jnp.ones((1, 1), dtype=int)))
Out[10]: 
{'loc': array([0]),
 'covariance_matrix': Array([[1]], dtype=int32),
 'precision_matrix': Array([[1.]], dtype=float32),
 'scale_tril': Array([[1.]], dtype=float32)}

Likewise, if we pass the precision matrix as int, it stays as int, but the cov matrix is float.

In [11]: get_args(dists.MultivariateNormal(0, precision_matrix=jnp.ones((1, 1), dtype=int)))
Out[11]: 
{'loc': array([0]),
 'covariance_matrix': Array([[1.]], dtype=float32),
 'precision_matrix': Array([[1]], dtype=int32),
 'scale_tril': Array([[1.]], dtype=float32)}

I think we can derive the dtype unambiguously from the constraint?