Realias result on the DiffResults paths (fixes #269 and #251)

Project.toml

@@ -1,6 +1,6 @@
 name = "ReverseDiff"
 uuid = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-version = "1.18.0-dev"
+version = "1.18.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/api/gradients.jl

@@ -21,25 +21,26 @@ call.
 function gradient(f, input, cfg::GradientConfig = GradientConfig(input))
     tape = GradientTape(f, input, cfg)
     result = construct_result(input_hook(tape))
-    seeded_reverse_pass!(result, tape)
+    result = seeded_reverse_pass!(result, tape)
     empty!(cfg.tape)
     return result
 end
 
 """
     ReverseDiff.gradient!(result, f, input, cfg::GradientConfig = GradientConfig(input))
 
-Returns `result`. This method is exactly like `ReverseDiff.gradient(f, input, cfg)`, except
-it stores the resulting gradient(s) in `result` rather than allocating new memory.
+This method is exactly like `ReverseDiff.gradient(f, input, cfg)`, except it stores the
+resulting gradient(s) in `result` rather than allocating new memory.
 
 `result` 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.
+in it as well. An immutable `DiffResult` cannot be updated in place and is replaced, so use
+the returned value: `result = ReverseDiff.gradient!(result, f, input, cfg)`.
 """
 function gradient!(result, f, input, cfg::GradientConfig = GradientConfig(input))
     tape = GradientTape(f, input, cfg)
-    seeded_reverse_pass!(result, tape)
+    result = seeded_reverse_pass!(result, tape)
     empty!(cfg.tape)
     return result
 end
@@ -60,23 +61,24 @@ of `f` w.r.t. `input[i].`
 """
 function gradient!(tape::Union{GradientTape,CompiledGradient}, input)
     result = construct_result(input_hook(tape))
-    gradient!(result, tape, input)
+    result = gradient!(result, tape, input)
     return result
 end
 
 """
     ReverseDiff.gradient!(result, tape::Union{GradientTape,CompiledGradient}, input)
 
-Returns `result`. This method is exactly like `ReverseDiff.gradient!(tape, input)`, except it
-stores the resulting gradient(s) in `result` rather than allocating new memory.
+This method is exactly like `ReverseDiff.gradient!(tape, input)`, except it stores the
+resulting gradient(s) in `result` rather than allocating new memory.
 
 `result` 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.
+in it as well. An immutable `DiffResult` cannot be updated in place and is replaced, so use
+the returned value: `result = ReverseDiff.gradient!(result, tape, input)`.
 """
 function gradient!(result, tape::Union{GradientTape,CompiledGradient}, input)
     seeded_forward_pass!(tape, input)
-    seeded_reverse_pass!(result, tape)
+    result = seeded_reverse_pass!(result, tape)
     return result
 end

src/api/hessians.jl

@@ -30,23 +30,25 @@ end
 
     ReverseDiff.hessian!(result::DiffResult, f, input::AbstractArray, cfg::HessianConfig = HessianConfig(result, input))
 
-Returns `result`. This method is exactly like `ReverseDiff.hessian(f, input, cfg)`, except
-it stores the resulting Hessian in `result` rather than allocating new memory.
+This method is exactly like `ReverseDiff.hessian(f, input, cfg)`, except it stores the
+resulting Hessian in `result` rather than allocating new memory.
 
 If `result` is a `DiffResults.DiffResult`, the primal value `f(input)` and the gradient
-`∇f(input)` will be stored in it along with the Hessian `H(f)(input)`.
+`∇f(input)` will be stored in it along with the Hessian `H(f)(input)`. An immutable
+`DiffResult` cannot be updated in place and is replaced, so use the returned value:
+`result = ReverseDiff.hessian!(result, f, input, cfg)`.
 """
 function hessian!(result, f, input::AbstractArray, cfg::HessianConfig = HessianConfig(input))
     ∇f = x -> gradient(f, x, cfg.gradient_config)
-    jacobian!(result, ∇f, input, cfg.jacobian_config)
+    result = jacobian!(result, ∇f, input, cfg.jacobian_config)
     return result
 end
 
 function hessian!(result::DiffResult, f, input::AbstractArray,
                   cfg::HessianConfig = HessianConfig(result, input))
     ∇f! = (y, x) -> begin
         gradient_result = DiffResult(zero(eltype(y)), y)
-        gradient!(gradient_result, f, x, cfg.gradient_config)
+        gradient_result = gradient!(gradient_result, f, x, cfg.gradient_config)
         result = DiffResults.value!(result, value(DiffResults.value(gradient_result)))
         return y
     end
@@ -68,7 +70,7 @@ return the Hessian `H(f)(input)`.
 """
 function hessian!(tape::Union{HessianTape,CompiledHessian}, input::AbstractArray)
     result = construct_result(output_hook(tape), input_hook(tape))
-    hessian!(result, tape, input)
+    result = hessian!(result, tape, input)
     return result
 end
 
@@ -77,21 +79,25 @@ end
 
     ReverseDiff.hessian!(result::DiffResult, tape::Union{HessianTape,CompiledHessian}, input)
 
-Returns `result`. This method is exactly like `ReverseDiff.hessian!(tape, input)`, except
-it stores the resulting Hessian in `result` rather than allocating new memory.
+This method is exactly like `ReverseDiff.hessian!(tape, input)`, except it stores the
+resulting Hessian in `result` rather than allocating new memory.
 
 If `result` is a `DiffResults.DiffResult`, the primal value `f(input)` and the gradient
-`∇f(input)` will be stored in it along with the Hessian `H(f)(input)`.
+`∇f(input)` will be stored in it along with the Hessian `H(f)(input)`. An immutable
+`DiffResult` cannot be updated in place and is replaced, so use the returned value:
+`result = ReverseDiff.hessian!(result, tape, input)`.
 """
 function hessian!(result::AbstractArray, tape::Union{HessianTape,CompiledHessian}, input::AbstractArray)
     seeded_forward_pass!(tape, input)
-    seeded_reverse_pass!(result, tape)
+    result = seeded_reverse_pass!(result, tape)
     return result
 end
 
 function hessian!(result::DiffResult, tape::Union{HessianTape,CompiledHessian}, input::AbstractArray)
     seeded_forward_pass!(tape, input)
-    seeded_reverse_pass!(DiffResult(DiffResults.gradient(result), DiffResults.hessian(result)), tape)
+    inner = DiffResult(DiffResults.gradient(result), DiffResults.hessian(result))
+    inner = seeded_reverse_pass!(inner, tape)
+    result = DiffResults.gradient!(result, DiffResults.value(inner))
     result = DiffResults.value!(result, func_hook(tape)(input))
     return result
 end

src/api/jacobians.jl

@@ -30,18 +30,19 @@ end
 """
     ReverseDiff.jacobian!(result, f, input, cfg::JacobianConfig = JacobianConfig(input))
 
-Returns `result`. This method is exactly like `ReverseDiff.jacobian(f, input, cfg)`, except
-it stores the resulting Jacobian(s) in `result` rather than allocating new memory.
+This method is exactly like `ReverseDiff.jacobian(f, input, cfg)`, except it stores the
+resulting Jacobian(s) in `result` rather than allocating new memory.
 
 `result` 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.
+in it as well. An immutable `DiffResult` cannot be updated in place and is replaced, so use
+the returned value: `result = ReverseDiff.jacobian!(result, f, input, cfg)`.
 """
 function jacobian!(result, f, input, cfg::JacobianConfig = JacobianConfig(input))
     tape = JacobianTape(f, input, cfg)
     isa(input, TrackedArray) && empty!(input.tape)
-    jacobian!(result, tape, input)
+    result = jacobian!(result, tape, input)
     empty!(tape.tape)
     return result
 end
@@ -74,7 +75,7 @@ form `f!(output::AbstractArray{<:Real}, input::AbstractArray{<:Real}...)`.
 function jacobian!(result, f!, output, input, cfg::JacobianConfig = JacobianConfig(output, input))
     tape = JacobianTape(f!, output, input, cfg)
     isa(input, TrackedArray) && empty!(input.tape)
-    jacobian!(result, tape, input)
+    result = jacobian!(result, tape, input)
     extract_result_value!(output, output_hook(tape))
     empty!(tape.tape)
     return result
@@ -103,23 +104,25 @@ new `output` values into the tape.
 """
 function jacobian!(tape::Union{JacobianTape,CompiledJacobian}, input)
     result = construct_result(output_hook(tape), input_hook(tape))
-    jacobian!(result, tape, input)
+    result = jacobian!(result, tape, input)
     return result
 end
 
 """
     ReverseDiff.jacobian!(result, tape::Union{JacobianTape,CompiledJacobian}, input)
 
-Returns `result`. This method is exactly like `ReverseDiff.jacobian!(tape, input)`, except it
-stores the resulting Jacobian(s) in `result` rather than allocating new memory.
+This method is exactly like `ReverseDiff.jacobian!(tape, input)`, except it stores the
+resulting Jacobian(s) in `result` rather than allocating new memory.
 
 `result` 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 of the target function will be stored in it as well.
+case the primal value of the target function will be stored in it as well. An immutable
+`DiffResult` cannot be updated in place and is replaced, so use the returned value:
+`result = ReverseDiff.jacobian!(result, tape, input)`.
 """
 function jacobian!(result, tape::Union{JacobianTape,CompiledJacobian}, input)
     seeded_forward_pass!(tape, input)
-    seeded_reverse_pass!(result, tape)
+    result = seeded_reverse_pass!(result, tape)
     return result
 end
 

src/api/tape.jl

@@ -44,7 +44,7 @@ function seeded_forward_pass!(t::AbstractTape, input)
 end
 
 function seeded_reverse_pass!(result, t::AbstractTape)
-    seeded_reverse_pass!(result, output_hook(t), input_hook(t), t)
+    result = seeded_reverse_pass!(result, output_hook(t), input_hook(t), t)
     return result
 end
 

src/api/utils.jl

@@ -29,12 +29,12 @@ function seeded_reverse_pass!(result, output::TrackedReal, input, tape)
     unseed!(input)
     seed!(output)
     reverse_pass!(tape)
-    extract_result!(result, output, input)
+    result = extract_result!(result, output, input)
     return result
 end
 
 function seeded_reverse_pass!(result, output::Number, input, tape)
-    extract_result!(result, output)
+    result = extract_result!(result, output, input)
     return result
 end
 
@@ -59,33 +59,20 @@ end
 
 function seeded_reverse_pass!(result::DiffResult, output::AbstractArray, input::TrackedArray, tape)
     seeded_reverse_pass!(DiffResults.jacobian(result), output, input, tape)
-    extract_result_value!(result, output)
+    result = extract_result_value!(result, output)
     return result
 end
 
-function seeded_reverse_pass!(result::Tuple, output::AbstractArray, input::Tuple, tape)
-    for i in eachindex(result)
-        seeded_reverse_pass!(result[i], output, input[i], tape)
-    end
-    return result
+function seeded_reverse_pass!(result::NTuple{N,Any}, output::AbstractArray, input::NTuple{N,Any}, tape) where {N}
+    return map((r, i) -> seeded_reverse_pass!(r, output, i, tape), result, input)
 end
 
 #####################
 # result extraction #
 #####################
 
-function extract_result!(result::Tuple, output, input::Tuple)
-    for i in eachindex(result)
-        extract_result!(result[i], output, input[i])
-    end
-    return result
-end
-
-function extract_result!(result::Tuple, output)
-    for i in eachindex(result)
-        extract_result!(result[i], output)
-    end
-    return result
+function extract_result!(result::NTuple{N,Any}, output, input::NTuple{N,Any}) where {N}
+    return map((r, i) -> extract_result!(r, output, i), result, input)
 end
 
 function extract_result!(result::AbstractArray, output::TrackedReal, input::TrackedArray)
@@ -99,24 +86,18 @@ function extract_result!(result::DiffResult, output::TrackedReal, input::Tracked
     return result
 end
 
-function extract_result!(result::AbstractArray, output::Number)
+# `input` is unused, but constrained as above so a mismatched `result` still fails
+function extract_result!(result::AbstractArray, output::Number, input::TrackedArray)
     fill_zeros!(result)
     return result
 end
 
-function extract_result!(result::DiffResult, output::Number)
+function extract_result!(result::DiffResult, output::Number, input::TrackedArray)
     result = DiffResults.value!(result, output)
     fill_zeros!(DiffResults.gradient(result))
     return result
 end
 
-function extract_result_value!(result::Tuple, output)
-    for i in eachindex(result)
-        extract_result_value!(result[i], output)
-    end
-    return result
-end
-
 function extract_result_value!(result::DiffResult, output::AbstractArray)
     result = DiffResults.value!(value, result, output)
     return result
@@ -132,11 +113,6 @@ function extract_result_value!(result::AbstractArray, output::AbstractArray)
     return result
 end
 
-function extract_result_value!(result::AbstractArray, output::TrackedArray)
-    copyto!(result, value(output))
-    return result
-end
-
 fill_zeros!(result::AbstractArray) = fill!(result, zero(eltype(result)))
 
 #######################