Skip to content

ENH: Add addional gufunc kernels for poisson_binom distribution#117

Open
steppi wants to merge 7 commits intoscipy:mainfrom
steppi:possion_binom_redux
Open

ENH: Add addional gufunc kernels for poisson_binom distribution#117
steppi wants to merge 7 commits intoscipy:mainfrom
steppi:possion_binom_redux

Conversation

@steppi
Copy link
Copy Markdown
Collaborator

@steppi steppi commented Mar 27, 2026

This PR adds a gufunc kernel to compute the cdf of the Poisson Binomial distribution, as well as reduction kernels to select elements from an array specifying values in a discrete distribution function, avoiding some Python overhead. I've incorporated these new kernels here, scipy/scipy#24894, to bring the performance of the proposed implementation of poisson.binom.pmf and poisson.binom.cdf in line with the performance of the existing implementation for small batch sizes.

@steppi
Copy link
Copy Markdown
Collaborator Author

steppi commented Mar 27, 2026

I've also changed all mdspans to use parens () for indexing instead of brackets [], to keep things in line with existing use of mdspans.

@steppi
Copy link
Copy Markdown
Collaborator Author

steppi commented Mar 27, 2026

@fbourgey, can you take a look at this when you get a chance?

@fbourgey
Copy link
Copy Markdown
Member

@fbourgey, can you take a look at this when you get a chance?

@steppi, of course. I'll review this in the coming days.

Copy link
Copy Markdown
Member

@fbourgey fbourgey left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@steppi LGTM, thanks!

  • Suggested one additional check.
  • Also, should we add a small test for poisson_binom_cdf_all?

Comment on lines +342 to +349
template <typename InputMat, typename OutputMat>
XSF_HOST_DEVICE inline void poisson_binom_cdf_all(InputMat p, OutputMat res) {
auto n = res.extent(0);
poisson_binom_pmf_all(p, res);
for (decltype(n) i = 1; i < n; i++) {
res(i) += res(i - 1);
}
}
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe add a specific error message (like in poisson_binom_pmf_all) when shapes are different?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants