Add MLX dispatch and a Metal kernel for betainc - #2368
Merged
Conversation
ricardoV94
reviewed
Aug 20, 2026
| np.testing.assert_allclose(left + right, 1.0, rtol=0.0, atol=1e-9) | ||
|
|
||
|
|
||
| def test_beta_logcdf(): |
Member
There was a problem hiding this comment.
why are we testing logcdfs and the like in pytensor?
Member
Author
There was a problem hiding this comment.
because it's the only real use-case for an integration test of this function. Happy to remove it.
Member
There was a problem hiding this comment.
the use case is we have a betainc that behaves like a reference backend. This seems more like a minimal pymc test, not a pytensor one?
They assert PyMC's logcdf expressions rather than that betainc matches the reference backend, which is what this suite is for.
ricardoV94
approved these changes
Aug 23, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
betaincwas the last op keepinglogcdfoff the MLX backend -- Beta, StudentT, SkewStudentT, Binomial and NegativeBinomial all go through it.The Metal kernel isn't a speedup layered on the vectorized dispatch; that dispatch is memory-bound. MLX materializes a buffer per iteration when a loop's coefficients are arrays, so 120 unrolled continued-fraction steps peak near 2.9 KB per element, which is 20 GiB at n=1e7, while the kernel keeps the recurrence in registers. It stays in as the float64 route and as the reference the kernel is tested against, since Metal is float32-only. Accuracy is 1e-15 to 1e-11 at float64 across the parameter plane, except where min(a, b) is large and a/(a+b) is small, which degrades to about 1e-5.