Closed tblazina closed 3 years ago
Thank you for taking this up! I will also need exhaustive tests before merging this PR (take the Beta distribution as an example). What we want to test for:
-inf
is returned by the logpdf when we're out of support.Of course! I plan to write tests, I just pushed this first commit to indicate that I'm working on it, I'll remove the [WIP]
whenever I think it's actually ready for review, and I'd be happy to get feedback on any numerical correctness! 😃
I see that the jax.scipy.stats.betabinom
code was just added last week and I was wondering if it makes sense to hold off on this PR until jax 0.2.10 is released rather than duplicating/implementing it here?
Yes, but JAX has a relatively short release cycle so we should be able to merge quickly; we would have to bump the JAX version manually and check that nothing breaks first.
Ok sounds good, I'll continue writing the tests and everything and then when JAX 0.2.10 is released and we see it hasn't broke anything, it can be merged. Should the update of jax happen in this PR or separately, do you have a preference?
Sounds good! Separate would be better.
Ok great, I'll finish this up, let you know when its ready for review and tag that it needs to be merged only after the PR for updating JAX is merged. 👍
Is it ready for review?
Sure! Other than the commented out tests which look at the logpdf correctness which can be tested once jax 0.2.10 is released. I more or less followed the tests in beta_tests.py
, not sure if it would also make sense to add some 'edge case' tests as is done in the binomial_tests.py
since the way it is implemented now is reusing the _random_binomial
function defined in the Binomial
distribution code.
Sorry had a sick kid at home. I'll make the changes and let you know when it's ready to merge!
No problem, that's life and we're in no hurry! Is the kid feeling better?
Ya she just started daycare so of course she got sick after 1 day and got us all sick. 🤷♂️ And my partner is applying for some professorships so has been very busy with application writing so ... child care has been falling on me until she sends those in. I'm full of excuses for why I have no time 😉
In any case I found a bit of time this morning - I added the changes you requested, they all seem good to me. The tests fail because the PDF correctness tests reply on that jax.scipy.stats.betabinomial
class which hasn't been released yet.
I just bumped JAX's version to 0.2.10, you can rebase to see if the tests pass.
Ok , updated the shape tests and with jax 0.2.10 the tests pass 👍 Unless yo have any other inputs, I think it's ready to merge.
Looks great to me! Thank you for taking the time to add this to the library!🙂
Add the BetaBinomial distribution