Replace incomplete_beta with betainc and speedup logcdf tests #4736
Conversation
|
Failing test was related to the minimum scipy version 1.2.0 not including I raised the requirement to 1.4.1, is this okay? There were some concerns about raising it to 1.6.0, because Colab is still on 1.4.1, but this should be fine in that regard. |
|
Looking at the pytest workflow, this PR seems to shave about 20-25 minutes from ~1h05m to 40m |
There was a problem hiding this comment.
This is the wrong place for an Aesara implementation of a scipy.special Op. Such Ops should be submitted as a PR to Aesara, so they can go alongside all the other scipy.special Ops; otherwise, we'll have more of the same unnecessary overlap in testing and maintenance that we've been trying to avoid.
Regardless, this is a very exciting update!
|
I can open a PR for the betainc in Aesara. When I opened the PR for the first time nobody mentioned that interest :) |
pymc3/distributions/dist_math.py
Outdated
| else: | ||
| dK = np.log1p(-x) + scipy.special.digamma(p + q) - scipy.special.digamma(q) | ||
|
|
||
| for n in range(1, max_iters + 1): |
There was a problem hiding this comment.
Is the plain python for loop going to be slow here, or is the function is meant to run in python/numpy/scipy
There was a problem hiding this comment.
Main Python loop for now. It's not really possible to vectorize as it is an iterative loop. It's still about ~4-5 orders of magnitude faster than the old C derivatives obtained by autodiff from the scan implementation of the incomplete_beta.
It's also very simple code to port to C or to Numba but I didn't want to do it just yet.
I have some speed comparisons here https://www.github.com/ricardoV94/derivatives_betainc/tree/master/comparison_aesara.ipynb
Sorry, I've been too busy to keep up on PRs, but, if you request a review, I'll usually get to it by the weekend. |
I didn't mean anyone was at fault, there is no rush on my side. I'll open it there soon. 😅 |
Also reverts reduced test n_samples due to speed issues
…nverseGamma` Closes pymc-devs#4467
|
Closing this one and opening a new one, since many things have changed |
Implement Betainc
This PR replaces the incomplete_beta method in dist_math by an aesara op
betaincthat wraps thescipy.special.betaincmethod. Supersedes #4519; Closes #4420.This provides a large increase in speed (and graph compilation) even without providing the c_code implementation directly (this can always be added later). I did some comparisons here. On my machine, the compilation time of a simple input/output function for the
incomplete_betawas around 10 seconds, and it's pretty much instantaneous for thebetaincas would be expected. Run time is ~2 orders of magnitude faster for the main op (even though it's calling scipy via Python), and ~4-5 orders of magnitude faster for the gradient.It also allows the evaluation of the logcdf for multiple values in Beta, StudentT, Binomial, NegativeBinomial (and zero inflated variants). I removed the previous limitation warning when trying to evaluate these methods for multiple values. These changes are critical for any practical implementation of Truncated Versions of these distributions.
The derivative algorithm was adapted from this paper and requires the use of two extra "pseudo aesara ops" that simply wrap the equivalent python function. This does not allow for further graph optimizations, but I doubt aesara could do much with the previous scan implementation either. I did some extensive comparison with the STAN implementation, and these derivatives seem to be both faster and more accurate than theirs, although it's not easy to find a good numerical reference. Some discussion can be found in stan-dev/math#2417
The one downside is the constant warning that the op has no
c-implementationwhenever the logp/logcdf makes use of it.Speedup logcdf tests
The logcdf tests are refactored to avoid recreating the logcdf graph for every parameter / value evaluation. This supersedes #4734. The old
incomplete_betawas found to be defective when attempting #4734, and so I decided to merge the two PRs into this one instead.This allows to remove most of the custom
n_sampleslimitations, and makes the test suite run considerably faster.A temporary work-around for the
HyperGeometrictests was added, related to an optimization issue on Aesara side pymc-devs/pytensor#461Finally, closes #4467