Add easy_rule for BigFloat division (#2934)

* Add easy_rule for BigFloat division

* Typo fix

* Fix?

* Add in multiplication too

* More rules and tests and mixed types

* Do inv too

* Add simple trig functions

* Fixed tests w autodiff

* Don't use rand

* Fix runtests.jl

---------

Co-authored-by: Katharine Hyatt <[email protected]>

Author: kshyatt

src/internal_rules/bigfloat.jl

@@ -7,13 +7,13 @@ function EnzymeRules.forward(
 
     if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
         if EnzymeRules.width(config) == 1
-            return RT(Ty.val(; kwargs...), Ty.val(; kwargs...))
+            return remove_innerty(RT)(Ty.val(; kwargs...), Ty.val(; kwargs...))
         else
             tup = ntuple(Val(EnzymeRules.width(config))) do i
                 Base.@_inline_meta
                 Ty.val(; kwargs...)
             end
-            return RT(Ty.val(; kwargs...), tup)
+            return remove_innerty(RT)(Ty.val(; kwargs...), tup)
         end
     elseif EnzymeRules.needs_shadow(config)
         if EnzymeRules.width(config) == 1
@@ -67,5 +67,19 @@ function EnzymeRules.reverse(
     return ()
 end
 
+EnzymeRules.@easy_rule(+(a::BigFloat, b::Number), (1,1))
+EnzymeRules.@easy_rule(+(a::Number, b::BigFloat), (1,1))
 EnzymeRules.@easy_rule(+(a::BigFloat, b::BigFloat), (1,1))
+EnzymeRules.@easy_rule(-(a::BigFloat, b::Number), (1,-1))
+EnzymeRules.@easy_rule(-(a::Number, b::BigFloat), (1,-1))
 EnzymeRules.@easy_rule(-(a::BigFloat, b::BigFloat), (1,-1))
+EnzymeRules.@easy_rule(*(a::BigFloat, b::BigFloat), (b, a))
+EnzymeRules.@easy_rule(*(a::BigFloat, b::Number), (b, a))
+EnzymeRules.@easy_rule(*(a::Number, b::BigFloat), (b, a))
+EnzymeRules.@easy_rule(/(a::BigFloat, b::Number), (one(a)/b, -(a/b^2)))
+EnzymeRules.@easy_rule(/(a::Number, b::BigFloat), (one(a)/b, -(a/b^2)))
+EnzymeRules.@easy_rule(/(a::BigFloat, b::BigFloat), (one(a)/b, -(a/b^2)))
+EnzymeRules.@easy_rule(Base.inv(a::BigFloat), (-(one(a)/a^2),))
+EnzymeRules.@easy_rule(Base.sin(a::BigFloat), (cos(a),))
+EnzymeRules.@easy_rule(Base.cos(a::BigFloat), (-sin(a),))
+EnzymeRules.@easy_rule(Base.tan(a::BigFloat), (one(a) + Ω^2,))

test/rules/internal_rules/bigfloat.jl

@@ -3,18 +3,27 @@ using EnzymeTestUtils
 using FiniteDifferences
 using Test
 
-@testset "BigFloat +/-" begin
-    a = rand(BigFloat)
-    b = rand(BigFloat)
+@testset "BigFloat arithmetic" begin
+    a = BigFloat(1.234)
+    da = BigFloat(-0.23)
+    b = BigFloat(0.56)
+    db = BigFloat(0.27)
+    af64 = 1.234 # for testing mixed methods
+    daf64 = -0.23 # for testing mixed methods
+    bf64 = 0.56 # for testing mixed methods
+    dbf64 = 0.27 # for testing mixed methods
 
-    # doesn't work because of https://github.com/EnzymeAD/Enzyme.jl/issues/2888
-    #test_reverse(+, Const, (a, Const), (b, Const))
-    #test_reverse(+, Active, (a, Active), (b, Active))
-    #test_reverse(-, Const, (a, Const), (b, Const))
-    #test_reverse(-, Active, (a, Active), (b, Active))
+    @test autodiff(Enzyme.Forward, +, Duplicated, Duplicated(a, da), Duplicated(b, db))[:1] ≈ da+db 
+    @test autodiff(Enzyme.Forward, +, Duplicated, Duplicated(a, da), Duplicated(bf64, dbf64))[:1] ≈ da+dbf64 
+    @test autodiff(Enzyme.Forward, -, Duplicated, Duplicated(a, da), Duplicated(b, db))[:1] ≈ da-db 
+    @test autodiff(Enzyme.Forward, -, Duplicated, Duplicated(a, da), Duplicated(bf64, dbf64))[:1] ≈ da-dbf64 
+    @test autodiff(Enzyme.Forward, *, Duplicated, Duplicated(a, da), Duplicated(b, db))[:1] ≈ b*da + a*db 
+    @test autodiff(Enzyme.Forward, *, Duplicated, Duplicated(a, da), Duplicated(bf64, dbf64))[:1] ≈ bf64*da + a*dbf64
+    @test autodiff(Enzyme.Forward, /, Duplicated, Duplicated(a, da), Duplicated(b, db))[:1] ≈ da/b  - db * a/b^2
+    @test autodiff(Enzyme.Forward, /, Duplicated, Duplicated(a, da), Duplicated(bf64, dbf64))[:1] ≈ da/bf64 - dbf64 * a/bf64^2 
 
-    test_forward(+, Const, (a, Const), (b, Const))
-    test_forward(+, Duplicated, (a, Duplicated), (b, Duplicated))
-    test_forward(-, Const, (a, Const), (b, Const))
-    test_forward(-, Duplicated, (a, Duplicated), (b, Duplicated))
+    @test autodiff(Enzyme.Forward, inv, Duplicated, Duplicated(a, da))[:1] ≈ -(one(BigFloat)/a^2) * da
+    @test autodiff(Enzyme.Forward, sin, Duplicated, Duplicated(a, da))[:1] ≈ cos(a) * da
+    @test autodiff(Enzyme.Forward, cos, Duplicated, Duplicated(a, da))[:1] ≈ -sin(a) * da 
+    @test autodiff(Enzyme.Forward, tan, Duplicated, Duplicated(a, da))[:1] ≈ autodiff(Enzyme.Forward, tan, Duplicated, Duplicated(af64, daf64))[1]
 end