[Feature Request] To return function evaluation results along with the computed gradient about reversediff.jl HOT 3 CLOSED

commit f22039512e6cbd10f1f19bb4a6368f87fcbab2ba
Author: ray@omni <[email protected]>
Date:   Sun Mar 13 19:20:24 2022 -0400

    Added `withgradient!`.

diff --git a/src/api/gradients.jl b/src/api/gradients.jl
index ed92cd4..ac58027 100644
--- a/src/api/gradients.jl
+++ b/src/api/gradients.jl
@@ -80,3 +80,39 @@ function gradient!(result, tape::Union{GradientTape,CompiledGradient}, input)
     seeded_reverse_pass!(result, tape)
     return result
 end
+
+"""
+    ReverseDiff.withgradient!(tape::Union{GradientTape,CompiledGradient}, input)
+
+If `input` is an `AbstractArray`, assume `tape` represents a function of the form
+`f(::AbstractArray)::Real` and return `∇f(input)`.
+
+If `input` is a tuple of `AbstractArray`s, assume `tape` represents a function of the form
+`f(::AbstractArray...)::Real` and return a `Tuple` where the `i`th element is the gradient
+of `f` w.r.t. `input[i].`
+"""
+function withgradient!(tape::Union{GradientTape,CompiledGradient}, input)
+    result = construct_result(input_hook(tape))
+    tmp = copy(value(output_hook(tape)))
+    output = isa(tmp, Number) ? [tmp] : copy(tmp)
+    withgradient!(output, result, tape, input)
+    return output, result
+end
+
+"""
+    ReverseDiff.withgradient!(output, result, tape::Union{GradientTape,CompiledGradient}, input)
+
+Returns `result`, and `output`. This method is like `ReverseDiff.gradient!(tape, input)`, except it
+stores the resulting gradient(s) in `result` rather than allocating new memory.
+
+`result` and `output` can be an `AbstractArray` or a `Tuple` of `AbstractArray`s. The `result` (or any
+of its elements, if `isa(result, Tuple)`), can also be a `DiffResults.DiffResult`, in which
+case the primal value `f(input)` (or `f(input...)`, if `isa(input, Tuple)`) will be stored
+in it as well.
+"""
+function withgradient!(output, result, tape::Union{GradientTape,CompiledGradient}, input)
+    seeded_forward_pass!(tape, input)
+    seeded_reverse_pass!(result, tape)
+    copyto!(output, value(output_hook(tape)))
+    return output, result
+end