WIP: add Enzyme support for fastpow

ext/DiffEqBaseEnzymeExt.jl

@@ -1,7 +1,7 @@
 module DiffEqBaseEnzymeExt
 
 using DiffEqBase
-import DiffEqBase: value
+import DiffEqBase: value, fastpow
 using Enzyme
 import Enzyme: Const
 using ChainRulesCore
@@ -53,4 +53,46 @@ function Enzyme.EnzymeRules.reverse(config::Enzyme.EnzymeRules.ConfigWidth{1},
     return ntuple(_ -> nothing, Val(length(args) + 4))
 end
 
+function Enzyme.EnzymeRules.forward(func::Const{typeof(DiffEqBase.fastpow)},
+        RT::Type{<:Union{Duplicated, DuplicatedNoNeed}},
+        _x::Union{Const, Duplicated}, _y::Union{Const, Duplicated})
+    x = _x.val
+    y = _y.val
+    ret = func.val(x, y)
+    if !(_x isa Const) 
+        dxval = _x.dval * y * (fastpow(x,y - 1))  
+    else 
+        dxval = make_zero(_x.val)
+    end
+    if !(_y isa Const)  
+        dyval = x isa Real && x<=0 ? Base.oftype(float(x), NaN) : _y.dval*(fastpow(x,y))*log(x)
+    else 
+        dyval = make_zero(_y.val)
+    end 
+    if RT <: DuplicatedNoNeed
+        return Float32(dxval + dyval)
+    else
+        return Duplicated(ret, Float32(dxval + dyval))
+    end
+end
+
+function EnzymeRules.augmented_primal(config::Enzyme.EnzymeRules.ConfigWidth{1}, 
+        func::Const{typeof(fastpow)}, ::Type{<:Active}, x::Active, y::Active)
+    if EnzymeRules.needs_primal(config)
+        primal = func.val(x.val, y.val)
+    else
+        primal = nothing
+    end
+    return EnzymeRules.AugmentedReturn(primal, nothing, nothing)
+end
+
+function EnzymeRules.reverse(config::Enzyme.EnzymeRules.ConfigWidth{1}, 
+        func::Const{typeof(DiffEqBase.fastpow)}, dret::Active, tape, _x::Active, _y::Active)
+    x = _x.val
+    y = _y.val
+    dxval =  dret.val * y * (fastpow(x,y - 1))
+    dyval = x isa Real && x<=0 ? Base.oftype(float(x), NaN) : dret.val * (fastpow(x,y))*log(x)
+    return (dxval, dyval)
 end
+
+end 

test/fastpow.jl

@@ -1,4 +1,5 @@
 using DiffEqBase: fastlog2, _exp2, fastpow
+using Enzyme, EnzymeTestUtils
 using Test
 
 @testset "Fast log2" begin
@@ -19,3 +20,23 @@ end
     errors = [abs(^(x, y) - fastpow(x, y)) for x in 0.001:0.001:1, y in 0.08:0.001:0.5]
     @test maximum(errors) < 1e-4
 end
+
+@testset "Fast pow - Enzyme forward rule" begin
+    @testset for RT in (Duplicated, DuplicatedNoNeed),
+        Tx in (Const, Duplicated),
+        Ty in (Const, Duplicated)
+        x = 3.0
+        y = 2.0
+        test_forward(fastpow, RT, (x, Tx), (y, Ty), atol=0.005, rtol=0.005)
+    end
+end
+
+@testset "Fast pow - Enzyme reverse rule" begin
+    @testset for RT in (Active,),
+        Tx in (Active,),
+        Ty in (Active,)
+        x = 2.0
+        y = 3.0
+        test_reverse(fastpow, RT, (x, Tx), (y, Ty), atol=0.001, rtol=0.001)
+    end
+end