Add Fastmath operators

src/lib/number.jl

@@ -1,4 +1,5 @@
 using DiffRules, SpecialFunctions, NaNMath
+using Base.FastMath: fast_op, make_fastmath
 
 @nograd isinf, isnan, isfinite, div
 
@@ -61,3 +62,31 @@ end
 @adjoint imag(x::Number) = imag(x), ī -> (real(ī)*im,)
 
 DiffRules._abs_deriv(x::Complex) = x/abs(x)
+
+ # Function add adjoint for Fastmath operations
+ # DiffRules is used to find derivatives for each operation
+for f in keys(fast_op)
+
+  fastf = fast_op[f]
+
+  if DiffRules.hasdiffrule(:Base, f, 1)
+    dx = DiffRules.diffrule(:Base, f, :x)
+    Δ = :Δ
+    if f in [:abs, :abs2]
+      Δ = :(real($Δ))
+    else
+      dx = :(conj($dx))
+    end
+    @eval begin
+      @adjoint Base.FastMath.$fastf(x::Number) =
+        Base.FastMath.$fastf(x), Δ -> ($Δ * make_fastmath($dx),)
+    end
+  elseif DiffRules.hasdiffrule(:Base, f, 2)
+    dx, dy = DiffRules.diffrule(:Base, f, :x, :y)
+    @eval begin
+      @adjoint Base.FastMath.$fastf(x::Number, y::Number) =
+        Base.FastMath.$fastf(x, y),
+        Δ -> (Δ * make_fastmath(conj($dx)), Δ * make_fastmath(conj($dy)))
+    end
+  end
+end

test/gradcheck.jl

@@ -983,3 +983,10 @@ end
   @test gradient(x -> findlast(ismissing, x), [1, missing]) == (nothing,)
   @test gradient(x -> findall(ismissing, x)[1], [1, missing]) == (nothing,)
 end
+
+@testset "fastmath" begin
+  @test gradient(x -> begin @fastmath sin(x) end, 1) == gradient(x -> sin(x), 1)
+  @test gradient(x -> begin @fastmath tanh(x) end, 1) == gradient(x -> tanh(x), 1)
+  @test gradient((x, y) -> begin @fastmath x*y end, 3, 2) == gradient((x, y) -> x*y, 3, 2)
+  @test gradient(x -> begin @fastmath real(log(x)) end, 1 + 2im) == gradient(x -> real(log(x)), 1 + 2im)
+end