Add EnzymeCore weakdep and an extension with a custom rule for the Levin transformation

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2021-2022 Michael Helton, Oscar Smith, and contributors
+Copyright (c) 2021-2023 Michael Helton, Oscar Smith, Chris Geoga, and contributors
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

Project.toml

@@ -6,11 +6,17 @@ version = "0.3.0-DEV"
 SIMDMath = "5443be0b-e40a-4f70-a07e-dcd652efc383"
 
 [compat]
-julia = "1.8"
 SIMDMath = "0.2.5"
+julia = "1.8"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+[weakdeps]
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+
+[extensions]
+BesselsEnzymeCoreExt = "EnzymeCore"
+
 [targets]
 test = ["Test"]

ext/BesselsEnzymeCoreExt.jl

@@ -0,0 +1,39 @@
+module BesselsEnzymeCoreExt
+
+  # TODO (cg 2023/05/08 10:02): Compat of any kind. 
+
+  using Bessels, EnzymeCore
+  using EnzymeCore.EnzymeRules
+  using Bessels.Math
+
+  # A manual method that takes an NTuple of partial sum terms and checks if it is
+  # exactly converged before computing the Levin sequence transformation. If it is
+  # exactly converged, the sequence transformation will create a divide-by-zero
+  # problem. See
+  #
+  # https://github.com/JuliaMath/Bessels.jl/issues/96 
+  #
+  # and links with for discussion.
+  #
+  # TODO (cg 2023/05/08 10:00): I'm not entirely sure how best to "generalize"
+  # this to cases like a return type of DuplicatedNoNeed, or something being a
+  # `Enzyme.Const`. These shouldn't in principle affect the "point" of this
+  # function (which is just to check for convergence before applying a
+  # function), but on its face this approach would mean I need a lot of
+  # hand-written extra methods. I have an open issue on the Enzyme.jl repo at
+  #
+  # https://github.com/EnzymeAD/Enzyme.jl/issues/786
+  #
+  # that gets at this problem a bit. But it's a weird request and I'm sure Billy
+  # has a lot of asks on his time.
+  function EnzymeRules.forward(func::Const{typeof(levin_transform)}, 
+                               ::Type{<:Duplicated}, 
+                               s::Duplicated,
+                               w::Duplicated)
+    (sv, dv, N) = (s.val, s.dval, length(s.val))
+    ls  = (sv[N-1] == sv[N]) ? sv[N] : levin_transform(sv, w.val)
+    dls = (dv[N-1] == dv[N]) ? dv[N] : levin_transform(dv, w.dval)
+    Duplicated(ls, dls)
+  end
+
+end

src/BesselFunctions/besselk.jl

@@ -578,15 +578,16 @@ end
 @generated function besselkx_levin(v, x::T, ::Val{N}) where {T <: FloatTypes, N}
     :(
         begin
-            s_0 = zero(T)
+            s = zero(T)
             t = one(T)
             @nexprs $N i -> begin
-                    s_{i} = s_{i-1} + t
-                    t *= (4*v^2 - (2i - 1)^2) / (8 * x * i)
-                    w_{i} = 1 / t
-                end
-                sequence = @ntuple $N i -> s_{i}
-                weights = @ntuple $N i -> w_{i}
+                  s += t
+                  t *= (4*v^2 - (2i - 1)^2) / (8 * x * i)
+                  s_{i} = s
+                  w_{i} = t
+              end
+              sequence = @ntuple $N i -> s_{i}
+              weights = @ntuple $N i -> w_{i}
             return levin_transform(sequence, weights) * sqrt(π / 2x)
         end
     )

src/Math/Math.jl

@@ -131,11 +131,12 @@ end
 #@inline levin_scale(B::T, n, k) where T = -(B + n) * (B + n + k)^(k - one(T)) / (B + n + k + one(T))^k
 @inline levin_scale(B::T, n, k) where T = -(B + n + k) * (B + n + k - 1) / ((B + n + 2k) * (B + n + 2k - 1))
 
-@inline @generated function levin_transform(s::NTuple{N, T}, w::NTuple{N, T}) where {N, T <: FloatTypes}
+@inline @generated function levin_transform(s::NTuple{N, T}, 
+                                            w::NTuple{N, T}) where {N, T <: FloatTypes}
     len = N - 1
     :(
         begin
-            @nexprs $N i -> a_{i} = Vec{2, T}((s[i] * w[i], w[i]))
+            @nexprs $N i -> a_{i} = Vec{2, T}((s[i] / w[i], 1 / w[i]))
             @nexprs $len k -> (@nexprs ($len-k) i -> a_{i} = fmadd(a_{i}, levin_scale(one(T), i, k-1), a_{i+1}))
             return (a_1[1] / a_1[2])
         end