Make sure to only fastmath the right things (#295)

* Make sure to only fastmath the right things

* spell non_transformed_definitions correctly

Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com>

Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com>

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.7.30"
+version = "0.7.31"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/base.jl

@@ -153,3 +153,8 @@ function rrule(::typeof(identity), x)
     return (x, identity_pullback)
 end
 
+# rouding related,
+# we use `zero` rather than `Zero()` for scalar, and avoids issues with map etc
+@scalar_rule round(x) zero(x)
+@scalar_rule floor(x) zero(x)
+@scalar_rule ceil(x) zero(x)

src/rulesets/Base/fastmath_able.jl

@@ -1,5 +1,8 @@
 let
     # Include inside this quote any rules that should have FastMath versions
+    # IMPORTANT:
+    # Do not add any rules here for functions that do not have varients in Base.FastMath
+    # e.g. do not add `foo` unless `Base.FastMath.foo_fast` exists.
     fastable_ast = quote
         #  Trig-Basics
         @scalar_rule cos(x) -(sin(x))
@@ -33,13 +36,6 @@ let
         @scalar_rule log10(x) inv(x) / log(oftype(x, 10))
         @scalar_rule log1p(x) inv(x + 1)
         @scalar_rule log2(x) inv(x) / log(oftype(x, 2))
-
-        # rouding related, 
-        # we use `zero` rather than `Zero()` for scalar, and avoids issues with map etc
-        @scalar_rule round(x) zero(x)
-        @scalar_rule floor(x) zero(x)
-        @scalar_rule ceil(x) zero(x)
-
 
         # Unary complex functions
         ## abs
@@ -195,10 +191,22 @@ let
             end
             return x * y, times_pullback
         end
-    end
+    end  # fastable_ast
 
     # Rewrite everything to use fast_math functions, including the type-constraints
-    eval(Base.FastMath.make_fastmath(fastable_ast))
+    fast_ast = Base.FastMath.make_fastmath(fastable_ast)
+
+    # Guard against mistakenly defining something as fast-able when it isn't.
+    non_transformed_definitions = intersect(fastable_ast.args, fast_ast.args)
+    filter!(expr->!(expr isa LineNumberNode), non_transformed_definitions)
+    if !isempty(non_transformed_definitions)
+        error(
+            "Non-FastMath compatible rules defined in fastmath_able.jl. \n Definitions:\n" *
+            join(non_transformed_definitions, "\n")
+        )
+    end
+
+    eval(fast_ast)
     eval(fastable_ast)  # Get original definitions
     # we do this second so it overwrites anything we included by mistake in the fastable
 end

test/rulesets/Base/base.jl

@@ -179,4 +179,17 @@
         frule_test(clamp, (x, Δx), (y, Δy), (z, Δz))
         rrule_test(clamp, Δk, (x, x̄), (y, ȳ), (z, z̄))
     end
+
+    @testset "rounding" begin
+        for x in (-0.6, -0.2, 0.1, 0.6)
+            # thanks to RoundNearest
+            if 0 > x % 1 > 0.5 || -1 < x % 1 <= 0.5
+                test_scalar(round, x; fdm=backward_fdm(5,1))
+            else
+                test_scalar(round, x; fdm=forward_fdm(5,1))
+            end
+            test_scalar(floor, x; fdm=backward_fdm(5, 1))
+            test_scalar(ceil, x; fdm=forward_fdm(5, 1))
+        end
+    end
 end

test/rulesets/Base/fastmath_able.jl

@@ -37,6 +37,9 @@ function complex_jacobian_test(f, z)
     @test jacobian_via_fdm(f, z) ≈ jacobian_via_rrule(f, z)
 end
 
+# IMPORTANT:
+# Do not add any tests here for functions that do not have varients in Base.FastMath
+# e.g. do not add `foo` unless `Base.FastMath.foo_fast` exists.
 const FASTABLE_AST = quote
     is_fastmath_mode = sin === Base.FastMath.sin_fast
 
@@ -90,19 +93,6 @@ const FASTABLE_AST = quote
         end
     end
 
-    @testset "rounding" begin
-        for x in (-0.6, -0.2, 0.1, 0.6)
-            # thanks to RoundNearest
-            if 0 > x % 1 > 0.5 || -1 < x % 1 <= 0.5
-                test_scalar(round, x; fdm=backward_fdm(5,1))
-            else
-                test_scalar(round, x; fdm=forward_fdm(5,1))
-            end
-            test_scalar(floor, x; fdm=backward_fdm(5, 1))
-            test_scalar(ceil, x; fdm=forward_fdm(5, 1))
-        end
-    end
-
     @testset "Unary complex functions" begin
         for f ∈ (abs, abs2, conj), z ∈ (-4.1-0.02im, 6.4, 3 + im)
             @testset "Unary complex functions f = $f, z = $z" begin