Add easy_rules for exponentially scaled Bessel functions (besselix etc.)#3365
Conversation
|
@wsmoses Can we approve CI here? |
|
The |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #3365 +/- ##
==========================================
- Coverage 77.00% 77.00% -0.01%
==========================================
Files 66 66
Lines 22285 22289 +4
==========================================
+ Hits 17161 17163 +2
- Misses 5124 5126 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
The test-job failures also look unrelated: they're The suite this PR actually adds to, |
Adds Enzyme rules for SpecialFunctions.besselix, besseljx, besselyx and besselkx (real arguments), which lower to AMOS routines (e.g. zbesi_) that Enzyme cannot differentiate through. Fixes EnzymeAD#2880 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TwgPLpFcdimYWqW6MVYQqK
Fixes #2880.
SpecialFunctions.besselix(and the other exponentially scaled Bessel functionsbesseljx,besselyx,besselkx) lower to AMOS routines (zbesi_, ...) that Enzyme cannot differentiate through, producingEnzymeNoDerivativeError: No augmented forward pass found for zbesi_.As suggested in the issue, this adds
EnzymeRules.@easy_ruledefinitions for all four scaled Bessel functions inEnzymeSpecialFunctionsExt, for real arguments, using the standard recurrence identities plus the derivative of the scaling factor:besselix(ν, x) = besseli(ν, x) * exp(-|x|)→(besselix(ν-1, x) + besselix(ν+1, x))/2 - sign(x)*Ωbesseljx(ν, x) = besselj(ν, x) * exp(-|imag(x)|)(scaling constant for realx) →(besseljx(ν-1, x) - besseljx(ν+1, x))/2besselyx(ν, x) = bessely(ν, x) * exp(-|imag(x)|)(scaling constant for realx) →(besselyx(ν-1, x) - besselyx(ν+1, x))/2besselkx(ν, x) = besselk(ν, x) * exp(x)→Ω - (besselkx(ν-1, x) + besselkx(ν+1, x))/2These match the corresponding ChainRules.jl rules. The order argument
νis marked@Constant(no order derivatives, same as the existingcmplx_known_opshandling of the unscaled Bessel functions).Tests: added
test_scalarcoverage for all four functions intest/ext/specialfunctions.jl. Verified locally that the issue reproducer now works and that reverse- and forward-mode derivatives match finite differences forν ∈ (0, 1, 2, 2.5)over positive/negative/zero real arguments (138 checks), including thesign(0) = 0branch ofbesselixatx = 0where the scaled function is only C¹ but the derivative is exact.🤖 Generated with Claude Code
https://claude.ai/code/session_01TwgPLpFcdimYWqW6MVYQqK