test: fix Enzyme test

Author: MilesCranmer

test/test_enzyme.jl

@@ -5,9 +5,6 @@ using DynamicExpressions
 operators = OperatorEnum(; binary_operators=(+, -, *, /), unary_operators=(cos, sin))
 # TODO: More operators will trigger a segfault in Enzyme
 
-# These options are required for Enzyme to work:
-const eval_options = (turbo=Val(false),)
-
 x1, x2, x3 = (i -> Node(Float64; feature=i)).(1:3)
 
 tree = Node(1, x1, Node(1, x2))  # == x1 + cos(x2)
@@ -16,32 +13,31 @@ X = randn(3, 100);
 dX = zero(X)
 
 function f(tree, X, operators, output)
-    output[] = sum(eval_tree_array(tree, X, operators; eval_options...)[1])
+    output[] = sum(eval_tree_array(tree, X, operators)[1])
     return nothing
 end
 
 output = [0.0]
 doutput = [1.0]
 
-autodiff(
-    Reverse,
-    f,
-    Const(tree),
-    Duplicated(X, dX),
-    Const(operators),
-    Duplicated(output, doutput),
-)
+fetch(schedule(Task(64 * 1024^2) do
+    autodiff(
+        Reverse,
+        f,
+        Const(tree),
+        Duplicated(X, dX),
+        Const(operators),
+        Duplicated(output, doutput),
+    )
+end))
 
 true_dX = cat(ones(100), -sin.(X[2, :]), zeros(100); dims=2)'
 
 @test true_dX ≈ dX
 
-#! format: off
-@static if false
-# Broken test (see https://github.com/EnzymeAD/Enzyme.jl/issues/1241)
 function my_loss_function(tree, X, operators)
     # Get the outputs
-    y = tree(X, operators)
+    y, _ = eval_tree_array(tree, X, operators)
     # Sum them (so we can take a gradient, rather than a jacobian)
     return sum(y)
 end
@@ -53,23 +49,18 @@ X = [1.0; 1.0;;]
 d_tree = begin
     storage_tree = copy(tree)
     # Set all constants to zero:
-    foreach(storage_tree) do node
-        if node.degree == 0 && node.constant
-            node.val = 0.0
-        end
-    end
-    autodiff(
-        Reverse,
-        my_loss_function,
-        Active,
-        Duplicated(tree, storage_tree),
-        Const(X),
-        Const(operators),
-    )
+    Enzyme.make_zero!(storage_tree)
+    fetch(schedule(Task(64 * 1024^2) do
+        autodiff(
+            Reverse,
+            my_loss_function,
+            Active,
+            Duplicated(tree, storage_tree),
+            Const(X),
+            Const(operators),
+        )
+    end))
     storage_tree
 end
 
-@test_broken get_scalar_constants(d_tree) ≈ [1.0, 0.717356]
-end
-
-#! format: on
+@test isapprox(first(get_scalar_constants(d_tree)), [1.0, 0.717356]; atol=1e-3)