Demo of convergence check with custom Enzyme method.

Author: Chris Geoga

ext/BesselsEnzymeCoreExt.jl

@@ -3,6 +3,7 @@ module BesselsEnzymeCoreExt
   using Bessels, EnzymeCore
   using EnzymeCore.EnzymeRules
   using Bessels.Math
+  import Bessels.Math: check_convergence
 
   # A manual method that separately transforms the `val` and `dval`, because
   # sometimes the `val` can converge while the `dval` hasn't, so just using an
@@ -17,18 +18,29 @@ module BesselsEnzymeCoreExt
       Duplicated(ls, dls)
   end
 
-  # This is fixing a straight bug in Enzyme.
+  function EnzymeRules.forward(func::Const{typeof(check_convergence)},
+                               ::Type{Const{Bool}},
+                               t::Duplicated{T}) where{T}
+    (t.val <= eps(T)) && (t.dval <= eps(T))
+  end
+
+  # This will be fixed upstream: see #861 for Enzyme.jl whenever the next
+  # release occurs.
   function EnzymeRules.forward(func::Const{typeof(sinpi)}, 
                                ::Type{<:Duplicated}, 
                                x::Duplicated)
       (sp, cp) = sincospi(x.val)
       Duplicated(sp, pi*cp*x.dval)
   end
 
+  # #861 will probably also mean this can be deleted at the next release of
+  # Enzyme.jl.
   function EnzymeRules.forward(func::Const{typeof(sinpi)}, 
                                ::Type{<:Const}, 
                                x::Const)
       sinpi(x.val)
   end
 
+
+
 end

src/BesselFunctions/besseli.jl

@@ -602,7 +602,7 @@ function besselix_large_args(v, x::ComplexOrReal{T}) where T
     for i in 1:MaxIter
         t *= -invx * ((4*v^2 - (2i - 1)^2) / i)
         s += t
-        abs(t) <= eps(T) && break
+        Math.check_convergence(t) && break
     end
     return s / sqrt(2 * (π * x))
 end

src/Math/Math.jl

@@ -154,7 +154,10 @@ end
     )
 end
 
-# TODO (cg 2023/05/16 18:09): dispute this cutoff.
 isnearint(x) = abs(x-round(x)) < 1e-5
 
+@inline check_convergence(term::T) where T = abs(term) <= eps(T)
+
+@inline check_convergence(ser::T, term::T) where T = (abs(term) <= eps(T)*ser)
+
 end