Add optional output_tangent kwarg to test_reverse (#2588)

Author: kshyatt

lib/EnzymeTestUtils/Project.toml

@@ -1,7 +1,7 @@
 name = "EnzymeTestUtils"
 uuid = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 authors = ["Seth Axen <[email protected]>", "William Moses <[email protected]>", "Valentin Churavy <[email protected]>"]
-version = "0.2.2"
+version = "0.2.3"
 
 [deps]
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"

lib/EnzymeTestUtils/src/test_reverse.jl

@@ -39,6 +39,7 @@ additional constraints:
 - `rtol`: Relative tolerance for `isapprox`.
 - `atol`: Absolute tolerance for `isapprox`.
 - `testset_name`: Name to use for a testset in which all tests are evaluated.
+- `output_tangent`: Optional final tangent to provide at the beginning of the reverse-mode differentiation 
 
 # Examples
 
@@ -76,8 +77,8 @@ function test_reverse(
     rtol::Real=1e-9,
     atol::Real=1e-9,
     testset_name=nothing,
-    runtime_activity::Bool=false
-)
+    runtime_activity::Bool=false,
+    output_tangent=nothing)
     call_with_captured_kwargs(f, xs...) = f(xs...; fkwargs...)
     if testset_name === nothing
         testset_name = "test_reverse: $f with return activity $ret_activity on $(_string_activity(args))"
@@ -92,12 +93,12 @@ function test_reverse(
         y = fcopy(args_copy...; deepcopy(fkwargs)...)
         # generate tangent for output
         if !_any_batch_duplicated(ret_activity, map(typeof, activities)...)
-            ȳ = ret_activity <: Const ? zero_tangent(y) : rand_tangent(rng, y)
+            ȳ = isnothing(output_tangent) ? (ret_activity <: Const ? zero_tangent(y) : rand_tangent(rng, y)) : output_tangent
         else
             batch_size = _batch_size(ret_activity, map(typeof, activities)...)
             ks = ntuple(Symbol ∘ string, batch_size)
             ȳ = ntuple(batch_size) do _
-                return ret_activity <: Const ? zero_tangent(y) : rand_tangent(y)
+                return isnothing(output_tangent) ? (ret_activity <: Const ? zero_tangent(y) : rand_tangent(rng, y)) : copy(output_tangent)
             end
         end
         # call finitedifferences, avoid mutating original arguments