Add new test mode rewind

[kshyatt]

This allows you to:

1. Test reverse mode as usual (against FD)
2. Run the AD in forward mode to generate output tangents
3. Use those output tangents to test reverse mode again (against FD) -- this implicitly tests the forward rule

This can be helpful when the tangents need a particular choice of gauge (e.g. the tangents for eigenvectors) -- FD will generate tangents with an arbitrary gauge.

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/lib/EnzymeTestUtils/src/test_rewind.jl b/lib/EnzymeTestUtils/src/test_rewind.jl
index 97649cc7..3f8653a2 100644
--- a/lib/EnzymeTestUtils/src/test_rewind.jl
+++ b/lib/EnzymeTestUtils/src/test_rewind.jl
@@ -61,26 +61,26 @@ end
 """
 
 function test_rewind(
-    f,
-    fwd_ret_activity,
-    rvs_ret_activity,
-    args...;
-    rng::Random.AbstractRNG=Random.default_rng(),
-    fdm=FiniteDifferences.central_fdm(5, 1),
-    fkwargs::NamedTuple=NamedTuple(),
-    rtol::Real=1e-9,
-    atol::Real=1e-9,
-    testset_name=nothing,
-    runtime_activity::Bool=false,
-    output_tangent=nothing,
-)
+        f,
+        fwd_ret_activity,
+        rvs_ret_activity,
+        args...;
+        rng::Random.AbstractRNG = Random.default_rng(),
+        fdm = FiniteDifferences.central_fdm(5, 1),
+        fkwargs::NamedTuple = NamedTuple(),
+        rtol::Real = 1.0e-9,
+        atol::Real = 1.0e-9,
+        testset_name = nothing,
+        runtime_activity::Bool = false,
+        output_tangent = nothing,
+    )
     # first, test reverse as normal with finite differences
-    test_reverse(f, rvs_ret_activity, args...; rng=rng, fdm=fdm, fkwargs=fkwargs, rtol=rtol, atol=atol, testset_name=testset_name, runtime_activity=runtime_activity, output_tangent=output_tangent)
-    # now, use the reverse rule to compare with the forward result 
+    test_reverse(f, rvs_ret_activity, args...; rng = rng, fdm = fdm, fkwargs = fkwargs, rtol = rtol, atol = atol, testset_name = testset_name, runtime_activity = runtime_activity, output_tangent = output_tangent)
+    # now, use the reverse rule to compare with the forward result
     if testset_name === nothing
         testset_name = "test_rewind: $f with return activity $fwd_ret_activity on $(_string_activity(args))"
     end
-    @testset "$testset_name" begin
+    return @testset "$testset_name" begin
         # test reverse rule to make sure it works with FD
         # run fwd mode first
 
@@ -109,6 +109,6 @@ function test_rewind(
         dy_ad = y_and_dy_ad[1]
         # now run this back through reverse mode, using dy_ad from forward mode
         # as the output tangent
-        test_reverse(f, rvs_ret_activity, args...; rng=rng, fdm=fdm, fkwargs=fkwargs, rtol=rtol, atol=atol, testset_name=testset_name, runtime_activity=runtime_activity, output_tangent=dy_ad)
+        test_reverse(f, rvs_ret_activity, args...; rng = rng, fdm = fdm, fkwargs = fkwargs, rtol = rtol, atol = atol, testset_name = testset_name, runtime_activity = runtime_activity, output_tangent = dy_ad)
     end
 end

[codecov]

Codecov Report

❌ Patch coverage is 0% with 23 lines in your changes missing coverage. Please review.
✅ Project coverage is 75.06%. Comparing base (3563bbd) to head (d4e4946).
⚠️ Report is 97 commits behind head on main.

Files with missing lines	Patch %	Lines
lib/EnzymeTestUtils/src/test_rewind.jl	0.00%	23 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2645      +/-   ##
==========================================
- Coverage   75.14%   75.06%   -0.09%     
==========================================
  Files          57       58       +1     
  Lines       17951    17974      +23     
==========================================
+ Hits        13490    13492       +2     
- Misses       4461     4482      +21     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[vchuravy]

Needs some tests/usage

Does it make sense to use fdm here? Or should we simply be testing for consistency between forward and reverse?

[vchuravy]

I must admit this is the first time I encounter the term gauge.

[kshyatt]

Ah sorry it's physicist brain

[kshyatt]

https://en.wikipedia.org/wiki/Gauge_fixing if you're interested. I'll try to rephrase

[kshyatt]

I used fdm here because the idea is to test at first that the reverse mode works at all, then test with the output tangents from forward mode. To put it another way, keeping the fdm allows for some "external" consistency check for my rules -- imagine I wrote a bad rule that just always returns 0 for all dvals. That would pass the consistency check of fwd/rvs -- but it would be wrong. So I think keeping the fdm sanity check is good.

[sethaxen]

In general I think forward-mode through an eigendecomposition will produce tangents "in the gauge" (I forget formally what the property is, something about the specific tangent space). But there's actually no guarantees in reverse-mode what kind of (co)-tangent you will have on the outputs. If downstream code interacts with the eigenvectors in a gauge-dependent manner, then your cotangents may not be "in the gauge" at all. So if you're implementing a rule, ideally you'd have one that works for all co-tangents; if that's not possible, then you'd need to require and document that downstream code interacts with the eigenvectors only in a gauge-invariant way. In that case, then I think this new testing method would work, but this seems highly specialized to this particular problem. If we think of the tangents as being constrained to a specific manifold, I can't actually remember off the top of my head if the manifold containing the co-tangent should always be a superset of this specific manifold.

It would be nice to have a second motivating application for this before adding another function to the API.

In terms of design, a few notes:

• For mutating functions test_forward and test_reverse will mutate the original arguments (including the function itself if it changes its own state), so you'll need some way to avoid this.
• Why is the first test_reverse unnecessary? Seems like it could just be left to the user to also run test_reverse themselves.
• I wonder then if this could just be an option for test_reverse, e.g. an option like output_tangent set to a singleton AutoFwd would indicate that forward-mode should be used to generate the output-tangent
• As an aside, I just saw that Add optional output_tangent kwarg to test_reverse #2588 added an output_tangent keyword. I think a more consistent API than adding the keyword output_tangent would be to allow ret_activity to be a tuple (output_tangent, activity), like the other arguments; then my above proposal would be to allow a user to provide e.g. (AutoFwd, Duplicated) to trigger this behavior.

[kshyatt]

Thanks for the comment! I agree this is pretty specialized -- it's just nearly all the functions I need to do AD with have this annoying gauge property. The problem isn't forward mode itself, it's the FD check (as FD generates tangents in a totally arbitrary gauge).

[kshyatt]

I think for now this can be closed, as after some thought our forward sensitivities are really bad to test against FD and I wasn't testing what I thought I was here...