test: simple test of chain rules

Author: MilesCranmer

Project.toml

@@ -30,6 +30,7 @@ DynamicExpressionsZygoteExt = "Zygote"
 [compat]
 Aqua = "0.7"
 Bumper = "0.6"
+ChainRulesCore = "1"
 Compat = "3.37, 4"
 Enzyme = "^0.11.12"
 LoopVectorization = "0.12"
@@ -51,11 +52,12 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "SafeTestsets", "Aqua", "Bumper", "Enzyme", "ForwardDiff", "LinearAlgebra", "LoopVectorization", "Optim", "SpecialFunctions", "StaticArrays", "SymbolicUtils", "Zygote"]
+test = ["Test", "SafeTestsets", "Aqua", "Bumper", "Enzyme", "ForwardDiff", "LinearAlgebra", "LoopVectorization", "Optim", "SpecialFunctions", "StaticArrays", "Suppressor", "SymbolicUtils", "Zygote"]

test/test_chainrules.jl

@@ -0,0 +1,54 @@
+using Test
+using DynamicExpressions
+using Random: MersenneTwister
+using ForwardDiff: gradient as fd_gradient
+using Zygote: gradient as zg_gradient
+using Suppressor: @suppress_err
+include("test_params.jl")
+include("tree_gen_utils.jl")
+
+let
+    rng = MersenneTwister(0)
+    n_features = 5
+    operators = OperatorEnum(; binary_operators=(+, *, -), unary_operators=(sin,))
+    tree = gen_random_tree_fixed_size(20, operators, n_features, Float64, Node, rng)
+    X = rand(rng, Float64, n_features, 100)
+
+    function f(X)
+        y, _ = eval_tree_array(tree, X, operators)
+        return sum(i -> y[i]^2, eachindex(y))
+    end
+
+    @suppress_err begin
+        # Check zg_gradient against fd_gradient; the latter of which is computed explicitly
+        @test isapprox([only(zg_gradient(f, X))...], [fd_gradient(f, X)...]; atol=1e-6)
+    end
+end
+
+mean(x) = sum(x) / length(x)
+
+let
+    operators = OperatorEnum(; binary_operators=(+, *, -), unary_operators=(sin,))
+    x1, x2, x3 = [Node{Float64}(; feature=i) for i in 1:3]
+    tree = sin(x1 * 3.2 - 0.9) + 0.2 * x2 - x3
+    X = [
+        1.0 2.0 3.0
+        4.0 5.0 6.0
+        7.0 8.0 9.0
+    ]
+    function eval_tree(X, tree)
+        y, _ = eval_tree_array(tree, X, operators)
+        return mean(y)
+    end
+
+    function true_eval_tree(X, c)
+        y = @. sin(X[1, :] * c[1] - c[2]) + c[3] * X[2, :] - X[3, :]
+        return mean(y)
+    end
+
+    evaluated_gradient = zg_gradient(tree -> eval_tree(X, tree), tree)[1]
+    true_gradient = fd_gradient(c -> true_eval_tree(X, c), [3.2, 0.9, 0.2])
+
+    @test evaluated_gradient.tree == tree
+    @test isapprox(evaluated_gradient.gradient, true_gradient)
+end

test/unittest.jl

@@ -31,6 +31,10 @@ end
     include("test_derivatives.jl")
 end
 
+@safetestset "Test chain rules" begin
+    include("test_chainrules.jl")
+end
+
 @safetestset "Test undefined derivatives" begin
     include("test_undefined_derivatives.jl")
 end