Exact chainrules derivatives for beta_inc and beta_inc_inv - #506
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #506 +/- ##
==========================================
+ Coverage 94.49% 95.03% +0.53%
==========================================
Files 14 14
Lines 3016 3200 +184
==========================================
+ Hits 2850 3041 +191
+ Misses 166 159 -7
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
@asinghvi17, @devmotion: I have switched implementations from the porting of the original code (which was probably not licensed correctly as we discussed yesterday) to a port of this repository, which is an independent implementation from @arzwa, unlicensed yet. I have better hopes that @arzwa will allow us to use his code, and have emailed him. By looking at the diff you'll see that this is clearly an independent implementation, and thus, if @arzwa agrees, we are good license-wise. I also have contacted the original authors, so if they come around faster we can still revert my last commit. |
|
I agree, you can use the code as you wish, I added an MIT license to the repo. |
|
I knew it'll be faster 🤣 |
Now that I'm reading mentioned discussion, I would like to add that I definitely implemented this using the approach described in the mentioned paper (Boik & Robinson-Cox 1998). As far as I remember, I did not 'translate' the code but implemented the described numerical method 'independently', based on their mathematical description (in fact I'm pretty sure I have never seen S-plus code in my life). However, I don't know what that would mean license-wise and whether it's OK to use an MIT license (I have never had to think about those kind of issues before). Any insights on this? |
|
I would test a much larger range of a,b values. Look at the code for calculation of the incomplete beta in this repo. They use different algorithms depending on the domain. |
devmotion
left a comment
devmotion
left a comment
There was a problem hiding this comment.
All the nested functions, multiple assignments on a single line and generally lack of comments make it quite challenging to follow the code. Maybe the implementation should be cleaned up a bit before it should be considered for inclusion in SpecialFunctions.
|
So, I have :
A slightly potentially-problematic outcome : the tests of the new chainrule alone last for ~40s on my machine, which is a large part of the total test time of the package :/ Ready for review round 2. |
|
Test time is probably not a concern. If this is good to go, we should merge. |
|
@devmotion Thanks a lot for the review and sorry for the delay ! All your suggestions were taken into account (I marked them as resolved), except 3: the first one is IMHO not a good idea for stability, the remaining two are points on which i agree with you but do not know how to do it myself, maybe you could help ? Tests passes, faillures on PRE seems unrelated. |
|
Thank you @lrnv for doing this. Pending its incorporation into SpecialFunctions, I have copied your code directly into a project I'm working on. I just wanted to check if this is a bug, in The asymmetry here looks funny, with |
|
The |
|
Sorry, I did not check the tests. Separately:
I assume the middle line is only needed because my copy of your code is outside of SpecialFunctions. I'm just using beta_inc() (actually Again, since I'm out of my depth, I leave it to you to judge if there's any practical upshot for you. |
|
Thanks @droodman for the hint on inlining, indeed it helped :) I am not sure the inclusion of ForwardDiff here is really in the scope of this PR. |
d7d26d4 to
500d94c
Compare
|
Hum... Apart ExplicitImports which yells about |
|
Hi @lrnv. Follow-up comments:
First the tolerance is set to eps(T)*T(1e4). Then it is lowered to 1e-14. I believe the original code uses 1e12. I tried changing to |
|
@droodman I removed the minapp/maxapp/eps weridness, thanks. I also removed the type restriction alltogether so that you can pass anything through. PS: I am not sure however that this is compliant with what SpecialFunctions.jl requires for merging ;) |
|
In the last couple of days, this code basically became the inner loop in my project and was consuming hours of time. So I was highly motivated to optimize it. I got another 2X improvement (at least on an M4 Mac), in addition (or multiplication) to the 2X improvement from adding @inline to functions. I forked @lrnv's fork of this repo and made a commit with the changes. However, I'm not that experienced with the complexities of GitHub--commits to different branches of different forks...--and I don't want to mess up anything that you are doing. Here is a link showing the new code block for ext/SpecialFunctionsChainRulesCoreExt.jl. As noted there, the main improvement is removing redundant calculations. I also deleted the first of the four return values of _beta_inc_grad() since it is never used. All tests pass. Timing code: Old timings: New timings: |
|
@droodman You created your branch out of my master branch, while you should have build it on top of my feature branch. Anyway sicne there is only one commit, I copy/pasted it here, thanks :) Ready again for review |
|
Hey @devmotion would you consider giving this another look ? |
d6071da to
7bc16eb
Compare
|
Rebased on master; faillure on pre seems unrelated. |
|
@devmotion Thanks for looking around ! I had troubles with Furthermore, as @droodman said, the numerous Still failling on pre for unrelated reasons. You may take another look. |
devmotion
left a comment
devmotion
left a comment
There was a problem hiding this comment.
Given the size of the diff I consulted Claude to make this review more thorough than I'd have managed by hand.
_beta_inc_grad agrees with central differences to 1e-10..1e-13 over a wide range of (a, b, x), including (500, 3, 0.99), (3, 500, 0.001) and (1e-3, 2, 0.01), at ~295ns vs ~230ns for the primal. So the continued fraction itself is not the problem — everything below is in the rules and the edge cases.
The 4-argument rule is off by a factor of 2
_beta_inc throws unless x + y ≈ 1 (src/beta_inc.jl:766), so y is not a free argument and there is no ∂/∂y at fixed x — perturbing one of them alone just leaves the domain. The only observable derivative is the one along the constraint. For f(z) = beta_inc(a, b, g(z), 1 - g(z))[1] the chain rule gives g'(z) * (∂x - ∂y), so whatever the rule does it has to satisfy
∂x - ∂y == ∂I/∂x
The PR returns (dIx, -dIx), which gives 2 * dIx. I checked: with g = logistic(z) and g = logistic(-z), a ∈ (0.4, 1.2, 45.0), b ∈ (2.3, 5.0, 16.0), z ∈ (-1.5, 0.0, 0.8), both frule and rrule come out at exactly 2.000x the finite difference in z — all 18 combinations.
Since x and y are a redundant parametrisation I'd split symmetrically, the same way we handle Symmetric/Hermitian cotangents:
(dIa, dIb, dIx / 2, -dIx / 2),
(-dIa, -dIb, -dIx / 2, dIx / 2),
With that all 18 cases give 1.000x. (dIx, ZeroTangent()) satisfies the identity as well but makes one argument a gradient sink, which seems worse in case most of the dependency happens to flow through y. Please don't use NoTangent() here — y is a perfectly differentiable Number and it would end up in a slot where ProjectTo expects a Real.
Worth getting right since beta_inc(a, b, x, y) is the form we advertise in docs/src/functions_overview.md:27, so it's what people will copy.
The current tests encode the factor 2 (test/chainrules.jl:316-318 and 340), so they need to change too. I'd drop the hand-written comparison of the individual partials and test the composition instead — that is testable with finite differences, unlike the raw 4-arg form, since z -> (g(z), 1 - g(z)) never leaves the domain:
@testset "4-arg beta_inc: derivative along the constraint y = 1 - x" begin
# ∂/∂x and ∂/∂y are not separately testable since beta_inc requires x + y == 1.
# What is observable is the derivative along the constraint, ie. for x = g(z),
# y = 1 - g(z) both modes have to reproduce
# d/dz I_{g(z)}(a, b) = pdf(g(z)) * g'(z),
# ie. the rule has to satisfy ∂x - ∂y == ∂I/∂x, however it splits the two.
logistic(z) = inv(1 + exp(-z))
# opposite signs of g' and x mirrored around 1/2, so that sign errors show up and
# the internal tail swap is exercised in both directions
gs = (
(name = "logistic(z)", g = z -> logistic(z), g′ = z -> logistic(z) * (1 - logistic(z))),
(name = "logistic(-z)", g = z -> logistic(-z), g′ = z -> -logistic(-z) * (1 - logistic(-z))),
)
for gg in gs, a in (0.4, 1.2, 45.0), b in (2.3, 5.0, 16.0), z in (-1.5, 0.0, 0.8)
x = gg.g(z)
d = gg.g′(z)
expected = x^(a - 1) * (1 - x)^(b - 1) / beta(a, b) * d
@testset "$(gg.name) a=$a b=$b z=$z" begin
@test all(isapprox.(beta_inc(a, b, x, 1 - x), beta_inc(a, b, x); rtol=1e-12))
_, Δ = frule((NoTangent(), 0.0, 0.0, d, -d), beta_inc, a, b, x, 1 - x)
@test isapprox(Δ[1], expected; rtol=1e-12)
@test isapprox(Δ[2], -expected; rtol=1e-12)
_, pb = rrule(beta_inc, a, b, x, 1 - x)
_, ā, b̄, x̄, ȳ = pb((1.0, 0.0))
@test isapprox(x̄ * d + ȳ * -d, expected; rtol=1e-12)
_, pb3 = rrule(beta_inc, a, b, x)
_, ā3, b̄3, _ = pb3((1.0, 0.0))
@test isapprox(ā, ā3; rtol=1e-12)
@test isapprox(b̄, b̄3; rtol=1e-12)
end
end
endAlso @test isapprox(ȳ4, -x̄3; rtol=1e-1) on line 337 — that relation is exact by construction and everything around it uses 1e-11, so I assume 1e-1 is left over from debugging.
_beta_inc_grad returns four values for x == 1
isone(x) && return oneT, zeroT, zeroT, zeroTEvery other branch returns three, and all three call sites destructure three. Julia silently drops the extra element, so dIa picks up the primal:
julia> ext._beta_inc_grad(3.0, 1.0, 1.0)
(1.0, 0.0, 0.0, 0.0)
julia> frule((NoTangent(), 1.0, 0.0, 0.0), beta_inc, 1.2, 2.3, 1.0)[2]
Tangent{Tuple{Float64, Float64}}(1.0, -1.0) # should be (0.0, 0.0)The x == 0 branch has the same arity problem, it just happens to be all zeros. Though ∂I/∂x at x == 0 is b for a == 1 and Inf for a < 1, not 0, so that branch isn't right either.
Independently of the wrong value this costs an allocation on every call, because the two arities widen the return type:
julia> Base.return_types(ext._beta_inc_grad, Tuple{Float64,Float64,Float64})
1-element Vector{Any}:
Union{Tuple{Float64, Float64, Float64}, NTuple{4, Float64}}
julia> @allocated ext._beta_inc_grad(1.2, 2.3, 0.4)
32We have allocation tests in the suite already, so an @inferred/@allocated check on _beta_inc_grad would be worth adding.
beta_inc_inv rule doesn't promote its arguments
The two beta_inc rules use map(float, promote(a, b, x)) but the beta_inc_inv one calls _beta_inc_grad(a, b, x) directly, so anything but Float64 arguments fails:
julia> frule((NoTangent(), 1.0, 0.0, 0.0), beta_inc_inv, 1, 2, 0.5)
ERROR: MethodError: no method matching _beta_inc_grad(::Int64, ::Int64, ::Float64)beta_inc_inv(1, 2, 0.5) itself works fine, so this is a regression introduced by the rule.
Float16 isn't supported although the primal supports it
_beta_inc_grad is restricted to Union{Float64, Float32}, but _beta_inc has a Union{Float16, Float32} method (src/beta_inc.jl:908) that computes in Float64 and converts back. Could you follow that pattern here? It fixes Float16, avoids compiling the whole CF a second time for Float32, and should let you tighten the rtol=5e-4 in the Float32 tests considerably.
Unconverged results are returned with just a @warn
maxapp = 200 isn't enough for large parameters, and the result can be quite wrong:
a, b, x |
rel. error of ∂I/∂a |
|
|---|---|---|
1000, 1000, 0.5 |
1.5e-12 | converged |
1e4, 1e4, 0.5 |
2.4e-10 | converged |
1e5, 1e5, 0.5 |
2.5e-8 | warns |
1e6, 1e6, 0.5 |
2.3e-2 | warns |
A 2% gradient isn't something a caller can be expected to notice. I'd either increase maxapp, or return NaN, or at the very least document the limitation. And the @warn needs maxlog=1 — as is it will be emitted on every gradient evaluation inside an optimiser or sampler. Note that @test_logs in test/chainrules.jl:290 pins the exact message, so that test would have to be adjusted along with it.
NaNs for degenerate arguments
julia> beta_inc(0.0, 2.0, 0.4), ext._beta_inc_grad(0.0, 2.0, 0.4)
((1.0, 0.0), (NaN, NaN, 0.0))
julia> frule((NoTangent(), 0.0, 0.0, 1.0), beta_inc_inv, 1.2, 2.3, 1.0)[2]
Tangent{Tuple{Float64, Float64}}(NaN, NaN)Same for b == 0 and for beta_inc_inv at p == 0. The beta_inc_inv ones come from 0 * inv(0) in -dIa * inv_dIx, partly aggravated by the x == 1 issue above. The primal handles all of these, so I think the rules should too, even if the answer is a deliberate NaN.
sign(Bn) in the guard against small Bn
invBn = (Bn > tiny || Bn < -tiny) && isfinite(Bn) ? inv(Bn) : inv(sign(Bn) * tiny)sign(0.0) == 0.0, so for Bn == 0 — the case the guard is there for — this returns Inf. inv(copysign(tiny, Bn)) would do what you want.
Smaller things
_dan_dqtakesqandda1_dqand uses neither, and its comment claims it returns the precomputed∂a₁/∂qforn == 1, which it doesn't. Same for then == 1remark in_dan_dp. The loop starts atn = 2, so I think all then == 1comments in these helpers can go, together with the unused arguments._dK_dqis missing thereturn, unlike_dK_dpright above it.An, an, Bn = _nextapp1(f, p, q)assigns the returnedbntoBn. It's correct sinceB₁ == b₁, but it reads like a bug — maybe just return(An, Bn)?ϵ = eps(T)*T(1e4)with the# 0) Previously keyword arguments:comment looks like leftover. Either make it a keyword likemaxapp/minappor drop the comment. Also the step numbering skips 2._dbn_dpand_dbn_dqrecomputeA,N1,NandDidentically and are both called every iteration.pfq_2(pf/q + 2) andpf_2q(pf + 2q) differing by one character is going to bite someone eventually.- Mixed
2*n/2p/2qspacing, and a double blank line between_dan_dpand_dan_dq. - The PR description cites Boik & Robinson-Cox 1998 (CSDA 27(1)) while the comment cites 1999 (JSS 3(1)). Both exist, but probably pick one. Keeping the reference in the comments is fine IMO, that's what we do in
src/beta_inc.jlas well. - The code says it's adapted from
arzwa/IncBetaDer.jl— could you check what licence that is under and whether a comment is enough?
Tests
The grid over (a, b, x) is good and it only adds ~30s, so no complaints there. But:
- Nothing near
x -> 0orx -> 1, which is exactly where the bug above is. - No
Integerarguments, noFloat16. test_frule/test_rruleare called without a surrounding@testset, so a failure gives you no idea which(a, b, x)it was. Could you wrap the loop bodies in@testset "a=$a b=$b x=$x"?
Other
beta_inc_inv(a, b, p, q) doesn't get a rule, although its 4-argument form is documented next to beta_inc(a, b, x, y) in the overview table. Could you add it with the same symmetric split?
Would you mind squashing/rewording the commits before merge?
Response to the latest review@devmotion below is a response to the various points you highlighted. The branch is now ready for another round of review and/or final merging, as you please.
I'm also using LLMs, no worries.
Agreed. I kept the continued-fraction algorithm and focused these changes on the rules, edge cases, convergence failure, type handling, and the internal cleanup listed below.
Fixed exactly as suggested. The four-argument rule now uses the symmetric
Updated. The tests now exercise the observable derivative along the constraint and expect one beta density, not twice the density. They also check the symmetric forward and reverse behavior across both tail-swap directions.
Fixed. The symmetric split is now checked with
Fixed. Every path returns exactly three derivatives. At both endpoints the shape derivatives are zero and the
Added both checks. The specialized
Fixed. The inverse rules now use
Done. The continued fraction is compiled only for
Changed to an adaptive but bounded safe mode. The loop exits as soon as all three values converge, with an empirical ceiling of 10,000 approximants instead of the historical JSS default of 200. This recovers difficult central cases that still produce accurate gradients (
Handled explicitly. For
Fixed at both division sites using
Fixed. These single-use helper formulas are now integrated directly into the
Added the explicit
Fixed.
Removed the leftover comment and numbered-step structure. The tolerance and iteration choices are documented as fixed
Fixed. Both derivatives are now computed together in the loop and share
Fixed. The abbreviated names were replaced with explicit names such as
Cleaned up the spacing and extra blank line. The two derivative formulas are now adjacent in their single loop call site.
The source consistently keeps the full 1999 JSS reference. I have updated the PR description to use the same reference and corrcted the details in it (w.r.t. the source and licences in particular).
The source comment now records the MIT licence and links to @arzwa's explicit permission to reuse the code: #506 (comment).
Added endpoint-limit tests for both
Added integer inverse-rule coverage and
Done. Both parameter grids now wrap each point in a named nested testset (using
Added the four-argument inverse rule with the same symmetric split between
Yes. I will squash and reword the branch history after these review changes are confirmed, before merge. Follow-up on
|
f5747af to
03540af
Compare
|
Thanks, this is a lot better. I went through every point again and they all hold up. As before I used Claude for the numerical checks, since most of these are only decidable by running them. Confirming the ones that mattered, at 03540af:
I also rederived
A few new things, one of which I would like fixed before merging.
|
| formulation | worst relative error |
|---|---|
| current | 4.75e-9, at a = 1e6, b = 1e8, x = 0.01 (1.4e-9 at a = b = 1e8) |
beta_integrand(a, b, x, 1 - x, -log(x) - log1p(-x)) |
1.31e-11 |
beta_integrand computes e^μ x^a y^b / B(a, b) with the DiDonato-Morris log1pmx/stirling_corr scheme that exists precisely to avoid this cancellation, and _beta_inc already relies on it. Its mu argument absorbs the x^-1 y^-1 factor without an intermediate underflow, which also gets subnormal cases like a = 45, b = 1e6, x = 1e-12 right.
SpecialFunctions.jl/src/beta_inc.jl
Lines 98 to 100 in 03540af
This matters more after this round than before it, since a = b = 1e8 is now a supported and tested case, and since the beta_inc_inv rules are built on inv(dIx) and inherit the loss.
It is not a strict improvement everywhere: at 55 of those 1813 points the beta_integrand route is more than 4x worse than the log form, although it never exceeds 1.31e-11 anywhere. So the question is whether a uniform 1.3e-11 bound is preferable to a 4.8e-9 tail. I think it is, and reusing what the primal uses is worth something on its own, but I do not feel strongly.
The a = b = 1e8 case is closer to the ceiling than the comment suggests
SpecialFunctions.jl/test/chainrules.jl
Lines 290 to 292 in 03540af
Bisecting, the last converging symmetric case is a = b ≈ 1.35e8; from about 1.5e8 upwards the rules return NaN while beta_inc itself is still perfectly happy at 1e10. So the test sits at roughly 64% of the budget and the cliff is only a factor 1.35 away.
The test is fine and NaN is much better than the old silent 2% error, so I am not asking for another increase. But of the three options I offered last time this leaves the third undone. Could you note the approximate ceiling in the comment above _beta_inc_grad? It will save whoever bisects this next.
SpecialFunctions.jl/ext/SpecialFunctionsChainRulesCoreExt.jl
Lines 372 to 374 in 03540af
Smaller things
One numbered step survived the cleanup:
SpecialFunctions.jl/ext/SpecialFunctionsChainRulesCoreExt.jl
Lines 587 to 589 in 03540af
The four-argument beta_inc rule promotes only three of the four arguments, so beta_inc(1.5f0, 2.25f0, 0.3f0, y::Float64) computes the derivative in Float32 against a Float64 primal. ProjectTo tidies up the tangent types so nothing breaks, but the four-argument inverse rule promotes all four, so this is at least inconsistent:
SpecialFunctions.jl/ext/SpecialFunctionsChainRulesCoreExt.jl
Lines 610 to 612 in 03540af
While you are there, _p in the three-argument inverse setup and _p, _q in the four-argument one are computed and never used:
SpecialFunctions.jl/ext/SpecialFunctionsChainRulesCoreExt.jl
Lines 633 to 635 in 03540af
Finally, on my earlier suggestion to tighten the Float32 tolerance: I was wrong that it could come down considerably. 1e-4 passes but 1e-5 does not, because ChainRulesTestUtils builds its finite-difference reference from the Float32 primal, so the reference is the limiting factor rather than the rule. Either tighten it to 1e-4 or leave it and say so in the comment, both are fine.
SpecialFunctions.jl/test/chainrules.jl
Lines 273 to 279 in 03540af
Only the NaN guard is blocking from my side. With that in, and the commits squashed, this is good to go.
Thank you for rechecking all of this, and especially for independently rederiving the recurrence and its partials. I kept the recurrence formulas unchanged and added a maintenance comment next to
Fixed.
Changed as suggested.
Documented. The comment above
Fixed: the remaining
Fixed. The four-argument
Fixed. The unused promoted values are now ignored explicitly in the destructuring assignments.
The
The blocking guard is now in place. I will squash the branch once again with good commit message. @devmotion, btw, did you note that i will talk about this PR at juliacon next week ? Come by if you want :) |
## Summary This PR adds ChainRules coverage for the regularized incomplete beta function and its inverse. - Adds analytic `frule`/`rrule` definitions for: - `beta_inc(a, b, x) -> (p, q)` - `beta_inc(a, b, x, y)`, with `x + y = 1` - `beta_inc_inv(a, b, p) -> (x, 1 - x)` - `beta_inc_inv(a, b, p, q)`, with `p + q = 1` - Uses symmetric tangent/cotangent splits for the redundant four-argument parametrizations. - Supports integer inputs through promotion and supports `Float16`/`Float32` by computing the differentiated continued fraction in `Float64` and converting back. - Handles endpoints and degenerate shape parameters explicitly. - Adds extensive tests for forward and reverse rules, endpoint behavior, type stability, allocations, promotion, non-convergence, and derivatives along the four-argument constraints. ## Implementation The parameter derivatives of `beta_inc` are computed by differentiating the continued-fraction method described by Boik & Robinson-Cox (1999). This is **not a translation of the S+ source code**: the implementation is independently adapted from [arzwa/IncBetaDer](https://github.com/arzwa/IncBetaDer), now MIT-licensed, whose author also [explicitly permitted reuse](JuliaMath#506 (comment)). The continued fraction exits as soon as its values converge. Its upper bound is `maxapp = 10_000`, rather than the historical value of 200, because difficult nearly symmetric cases with large shape parameters require substantially more approximants. If convergence is not reached, the helper emits a rate-limited warning and returns `NaN` derivatives instead of silently returning an inaccurate gradient. Small recurrence helpers used only once have been integrated into the loop. The remaining shared helper stays `@inline` because it is called four times per iteration and benchmarking showed a material performance benefit. ## Motivation This allows automatic differentiation through `beta_inc` and `beta_inc_inv`. Missing derivatives affect, among other use cases: - beta quantiles and beta-distribution marginals; - Student-t and F distribution workflows built on incomplete beta functions; - fitting `Distributions.MvTDist`; - applications in packages such as Copulas.jl; - the use case discussed in [this Julia Discourse thread](https://discourse.julialang.org/t/gaussian-copula-priors-in-turing-auto-differentiation-error-with-beta-quantile-function/132779/5). These rules naturally belong in SpecialFunctions.jl's ChainRulesCore extension. ## Validation The tests cover a broad grid of `(a, b, x)`, including values close to `x = 0` and `x = 1`, integer and reduced-precision inputs, and the constrained four-argument APIs. Sensitive large-parameter regression cases include `a = b = 10^6` and `a = b = 10^8` at `x = 0.5`; these exercise cases that need more than 200 and, for `10^8`, thousands of approximants. The full SpecialFunctions.jl test suite passes, and the new code is covered by the test suite. ## Reference > Boik, R. J., & Robinson-Cox, J. F. (1999). Derivatives of the incomplete beta function. *Journal of Statistical Software, 3*(1). https://doi.org/10.18637/jss.v003.i01
8afa458 to
988b671
Compare
## Summary This PR adds ChainRules coverage for the regularized incomplete beta function and its inverse. - Adds analytic `frule`/`rrule` definitions for: - `beta_inc(a, b, x) -> (p, q)` - `beta_inc(a, b, x, y)`, with `x + y = 1` - `beta_inc_inv(a, b, p) -> (x, 1 - x)` - `beta_inc_inv(a, b, p, q)`, with `p + q = 1` - Uses symmetric tangent/cotangent splits for the redundant four-argument parametrizations. - Supports integer inputs through promotion and supports `Float16`/`Float32` by computing the differentiated continued fraction in `Float64` and converting back. - Handles endpoints and degenerate shape parameters explicitly. - Adds extensive tests for forward and reverse rules, endpoint behavior, type stability, allocations, promotion, non-convergence, and derivatives along the four-argument constraints. ## Implementation The parameter derivatives of `beta_inc` are computed by differentiating the continued-fraction method described by Boik & Robinson-Cox (1999). This is **not a translation of the S+ source code**: the implementation is independently adapted from [arzwa/IncBetaDer](https://github.com/arzwa/IncBetaDer), now MIT-licensed, whose author also [explicitly permitted reuse](JuliaMath#506 (comment)). The continued fraction exits as soon as its values converge. Its upper bound is `maxapp = 10_000`, rather than the historical value of 200, because difficult nearly symmetric cases with large shape parameters require substantially more approximants. If convergence is not reached, the helper emits a rate-limited warning and returns `NaN` derivatives instead of silently returning an inaccurate gradient. Small recurrence helpers used only once have been integrated into the loop. The remaining shared helper stays `@inline` because it is called four times per iteration and benchmarking showed a material performance benefit. ## Motivation This allows automatic differentiation through `beta_inc` and `beta_inc_inv`. Missing derivatives affect, among other use cases: - beta quantiles and beta-distribution marginals; - Student-t and F distribution workflows built on incomplete beta functions; - fitting `Distributions.MvTDist`; - applications in packages such as Copulas.jl; - the use case discussed in [this Julia Discourse thread](https://discourse.julialang.org/t/gaussian-copula-priors-in-turing-auto-differentiation-error-with-beta-quantile-function/132779/5). These rules naturally belong in SpecialFunctions.jl's ChainRulesCore extension. ## Validation The tests cover a broad grid of `(a, b, x)`, including values close to `x = 0` and `x = 1`, integer and reduced-precision inputs, and the constrained four-argument APIs. Sensitive large-parameter regression cases include `a = b = 10^6` and `a = b = 10^8` at `x = 0.5`; these exercise cases that need more than 200 and, for `10^8`, thousands of approximants. The full SpecialFunctions.jl test suite passes, and the new code is covered by the test suite. ## Reference > Boik, R. J., & Robinson-Cox, J. F. (1999). Derivatives of the incomplete beta function. *Journal of Statistical Software, 3*(1). https://doi.org/10.18637/jss.v003.i01
988b671 to
2e3418f
Compare
|
Thanks, the blocking item is properly fixed. I went through the round again at 2e3418f, in the same way as before, and I also spent some time on the parts of the algorithm that neither of us had looked at closely yet. Confirming what changed:
I also checked whether Two things I got wrong last time, before the new items. The subnormal example was the wrong way roundI said SpecialFunctions.jl/test/chainrules.jl Lines 321 to 327 in 2e3418f Both are within two ulp of the subnormal grid, so nothing is broken, but the third point now enshrines the less accurate of the two values. The The 2e-9 figure was too optimistic
The mixed promotion test does not test the promotionSpecialFunctions.jl/test/chainrules.jl Lines 346 to 351 in 2e3418f The tangents are The fix itself is correct, only the test cannot see it. Asserting on the derivative tuple works: d = ChainRulesCore.derivatives_given_output(Ω, beta_inc, 1.5f0, 2.25f0, 0.3f0, 1 - Float64(0.3f0))
@test eltype(first(d)) === Float64The convergence test stops being a relative test below epsSpecialFunctions.jl/ext/SpecialFunctionsChainRulesCoreExt.jl Lines 566 to 572 in 2e3418f The
The affected values are all far below anything that matters in a gradient, and I checked that: over 68880 points, restricted to The documented ceiling only covers the symmetric caseSpecialFunctions.jl/ext/SpecialFunctionsChainRulesCoreExt.jl Lines 377 to 380 in 2e3418f Asymmetric shapes stop converging at roughly half the symmetric value:
Version
Smaller thingsThe With The remaining error at large shapes is not in the continued fraction, it is in the prefactor: SpecialFunctions.jl/ext/SpecialFunctionsChainRulesCoreExt.jl Lines 336 to 338 in 2e3418f The bracket Strongly imbalanced shapes return NaNs where the primal is fine. The four-argument On the tightened tolerance: Finally, the licence note names IncBetaDer.jl and links the permission, but does not reproduce the upstream copyright line. The precedent in this repo is |
## Summary This PR adds ChainRules coverage for the regularized incomplete beta function and its inverse. - Adds analytic `frule`/`rrule` definitions for: - `beta_inc(a, b, x) -> (p, q)` - `beta_inc(a, b, x, y)`, with `x + y = 1` - `beta_inc_inv(a, b, p) -> (x, 1 - x)` - `beta_inc_inv(a, b, p, q)`, with `p + q = 1` - Uses symmetric tangent/cotangent splits for the redundant four-argument parametrizations. - Supports integer inputs through promotion and supports `Float16`/`Float32` by computing the differentiated continued fraction in `Float64` and converting back. - Handles endpoints and degenerate shape parameters explicitly. - Adds extensive tests for forward and reverse rules, endpoint behavior, type stability, allocations, promotion, non-convergence, and derivatives along the four-argument constraints. ## Implementation The parameter derivatives of `beta_inc` are computed by differentiating the continued-fraction method described by Boik & Robinson-Cox (1999). This is **not a translation of the S+ source code**: the implementation is independently adapted from [arzwa/IncBetaDer](https://github.com/arzwa/IncBetaDer), now MIT-licensed, whose author also [explicitly permitted reuse](JuliaMath#506 (comment)). The continued fraction exits as soon as its values converge. Its upper bound is `maxapp = 10_000`, rather than the historical value of 200, because difficult nearly symmetric cases with large shape parameters require substantially more approximants. If convergence is not reached, the helper emits a rate-limited warning and returns `NaN` derivatives instead of silently returning an inaccurate gradient. Small recurrence helpers used only once have been integrated into the loop. The remaining shared helper stays `@inline` because it is called four times per iteration and benchmarking showed a material performance benefit. ## Motivation This allows automatic differentiation through `beta_inc` and `beta_inc_inv`. Missing derivatives affect, among other use cases: - beta quantiles and beta-distribution marginals; - Student-t and F distribution workflows built on incomplete beta functions; - fitting `Distributions.MvTDist`; - applications in packages such as Copulas.jl; - the use case discussed in [this Julia Discourse thread](https://discourse.julialang.org/t/gaussian-copula-priors-in-turing-auto-differentiation-error-with-beta-quantile-function/132779/5). These rules naturally belong in SpecialFunctions.jl's ChainRulesCore extension. ## Validation The tests cover a broad grid of `(a, b, x)`, including values close to `x = 0` and `x = 1`, integer and reduced-precision inputs, and the constrained four-argument APIs. Sensitive large-parameter regression cases include `a = b = 10^6` and `a = b = 10^8` at `x = 0.5`; these exercise cases that need more than 200 and, for `10^8`, thousands of approximants. The full SpecialFunctions.jl test suite passes, and the new code is covered by the test suite. ## Reference > Boik, R. J., & Robinson-Cox, J. F. (1999). Derivatives of the incomplete beta function. *Journal of Statistical Software, 3*(1). https://doi.org/10.18637/jss.v003.i01
2e3418f to
c91356f
Compare
|
Thank you for another very thorough numerical pass. I addressed the test and implementation issues below, and recorded the useful algorithmic limitations in the code where a change would require a separate numerical study.
Fixed. I dropped the subnormal point and replaced the two remaining self-comparisons with hard-coded 600-bit references at the asymmetric and symmetric large-shape points. The new tests fail with the old log formulation and can detect a regression in
Agreed; thank you for correcting the record. I no longer claim the 2e-9 figure. The remaining large-shape error is documented at the prefactor below.
Fixed exactly along these lines. The regression test now inspects
Fixed. The default is now
Documented. The function comment now gives approximately
Fixed: the package version is now
I kept
Agreed; I did not change the formula in this PR. I added a comment directly above the bracket documenting the cancellation near
I kept the 10,000 default deliberately. Raising it globally would make genuinely nonconverging calls several times more expensive, and this imbalanced regime is outside the paper's validated range; the explicit NaN is preferable to silently accepting an immature gradient. The function comment now documents
Fixed. The differentiated helper now has a four-argument path that carries the caller's
Fixed without relying on a favorable random draw. All four Float32 checks now use explicit, moderate input tangents, and both reverse-rule checks use an explicit non-symmetric output cotangent. They therefore no longer depend on Julia's RNG stream or on future changes to
Fixed. The source now includes the upstream line |
7 participants
Summary
This PR adds ChainRules coverage for the regularized incomplete beta function and its inverse.
frule/rruledefinitions for:beta_inc(a, b, x) -> (p, q)beta_inc(a, b, x, y), withx + y = 1beta_inc_inv(a, b, p) -> (x, 1 - x)beta_inc_inv(a, b, p, q), withp + q = 1Float16/Float32by computing the differentiated continued fraction inFloat64and converting back.Implementation
The parameter derivatives of
beta_incare computed by differentiating the continued-fraction method described by Boik & Robinson-Cox (1999). This is not a translation of the S+ source code: the implementation is independently adapted from arzwa/IncBetaDer, now MIT-licensed, whose author also explicitly permitted reuse.The continued fraction exits as soon as its values converge. Its upper bound is
maxapp = 10_000, rather than the historical value of 200, because difficult nearly symmetric cases with large shape parameters require substantially more approximants. If convergence is not reached, the helper emits a rate-limited warning and returnsNaNderivatives instead of silently returning an inaccurate gradient.Small recurrence helpers used only once have been integrated into the loop. The remaining shared helper stays
@inlinebecause it is called four times per iteration and benchmarking showed a material performance benefit.Motivation
This allows automatic differentiation through
beta_incandbeta_inc_inv. Missing derivatives affect, among other use cases:Distributions.MvTDist;These rules naturally belong in SpecialFunctions.jl's ChainRulesCore extension.
Validation
The tests cover a broad grid of
(a, b, x), including values close tox = 0andx = 1, integer and reduced-precision inputs, and the constrained four-argument APIs. Sensitive large-parameter regression cases includea = b = 10^6anda = b = 10^8atx = 0.5; these exercise cases that need more than 200 and, for10^8, thousands of approximants.The full SpecialFunctions.jl test suite passes, and the new code is covered by the test suite.
Reference