pyro-ppl / numpyro / 39573199749

93%

push

github

fix(distributions): sample binomials with integer-valued float total counts (#2382)

* fix(distributions): sample binomials with integer-valued float total counts

I found that Binomial and BetaBinomial sampling raised a TypeError when total_count was an integer-valued float array, although such counts pass the constraint and work in log_prob. The binomial sampler now counts trials with integers and returns the draws in the dtype of total_count, so float counts give the same draws as integer counts.

Assisted-by: Claude
Assisted-by: GPT-6 Astra
Signed-off-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com>

* test(distributions): parametrize the binomial test over the count dtype

I replaced the test with the suggested one over float32, int16 and uint8 counts, which also covers the non-default integer dtypes that failed on master, and widened the comment in _binomial to match.

Assisted-by: Claude
Assisted-by: GPT-6 Astra
Signed-off-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com>

---------

Signed-off-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com>
Co-authored-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com>

3 of 3 new or added lines in 1 file covered. (100.0%)

15796 of 17010 relevant lines covered (92.86%)

1.7 hits per line