diff --git a/Project.toml b/Project.toml index 596d5dd1..a24ae60d 100644 --- a/Project.toml +++ b/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" diff --git a/src/api/gradients.jl b/src/api/gradients.jl index ed92cd4d..234a44a8 100644 --- a/src/api/gradients.jl +++ b/src/api/gradients.jl @@ -21,7 +21,7 @@ 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 @@ -29,17 +29,18 @@ 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 diff --git a/src/api/hessians.jl b/src/api/hessians.jl index e0107cb9..db66fec2 100644 --- a/src/api/hessians.jl +++ b/src/api/hessians.jl @@ -30,15 +30,17 @@ 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 @@ -46,7 +48,7 @@ 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 diff --git a/src/api/jacobians.jl b/src/api/jacobians.jl index 132957af..58c9fc5a 100644 --- a/src/api/jacobians.jl +++ b/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 diff --git a/src/api/tape.jl b/src/api/tape.jl index 2f033ee6..1be1f973 100644 --- a/src/api/tape.jl +++ b/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 diff --git a/src/api/utils.jl b/src/api/utils.jl index a2469b0e..147cbd80 100644 --- a/src/api/utils.jl +++ b/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))) ####################### diff --git a/test/api/GradientTests.jl b/test/api/GradientTests.jl index f05868c9..27b26840 100644 --- a/test/api/GradientTests.jl +++ b/test/api/GradientTests.jl @@ -1,6 +1,6 @@ module GradientTests -using DiffTests, ForwardDiff, ReverseDiff, Test, LinearAlgebra +using DiffTests, ForwardDiff, ReverseDiff, StaticArrays, Test, LinearAlgebra include(joinpath(dirname(@__FILE__), "../utils.jl")) @@ -261,4 +261,109 @@ end end end +############################################################################################ + +# Top level, not inside the `@testset`: closures would change the inlining decisions. +f269(x) = sum(abs2, x) + +function value_and_gradient269!(grad, tape, x) + result = DiffResults.MutableDiffResult(zero(eltype(x)), (grad,)) + result = ReverseDiff.gradient!(result, tape, x) + return DiffResults.value(result), DiffResults.gradient(result) +end + +nested269!(grad, tape, x) = (y = value_and_gradient269!(grad, tape, x)[1]; (y, grad)) + +@testset "primal value survives inlining into a caller (#269)" begin + x = [3.0, 5.0] + tape = ReverseDiff.GradientTape(f269, x) + for t in (tape, ReverseDiff.compile(tape)) + @test nested269!(similar(x), t, x) == (34.0, [6.0, 10.0]) + end +end + +############################################################################################ + +f251(x) = sum(abs2, x) + +# An `MVector` gives an `ImmutableDiffResult` with a writable gradient buffer; an `SVector` +# buffer cannot be written to at all. +@testset "primal value of an immutable result (#251)" begin + x = MVector{2}(3.0, 5.0) + value, grad = 34.0, [6.0, 10.0] + + # `GradientResult` aliases its argument as the gradient buffer, so pass a copy. + result = ReverseDiff.gradient!(DiffResults.GradientResult(MVector(x)), f251, x) + @test result isa DiffResults.ImmutableDiffResult + @test DiffResults.value(result) == value + @test DiffResults.gradient(result) == grad + + tape = ReverseDiff.GradientTape(f251, x) + for t in (tape, ReverseDiff.compile(tape)) + result = ReverseDiff.gradient!(DiffResults.GradientResult(MVector(x)), t, x) + @test DiffResults.value(result) == value + @test DiffResults.gradient(result) == grad + end +end + +g251(x, y) = sum(abs2, x) + sum(abs2, y) + +@testset "primal value of immutable results in a tuple (#251)" begin + x, y = MVector{2}(3.0, 5.0), MVector{2}(2.0, 4.0) + value, grads = 54.0, ([6.0, 10.0], [4.0, 8.0]) + + result = (DiffResults.GradientResult(MVector(x)), DiffResults.GradientResult(MVector(y))) + result = ReverseDiff.gradient!(result, g251, (x, y)) + @test all(r -> r isa DiffResults.ImmutableDiffResult, result) + @test map(DiffResults.value, result) == (value, value) + @test map(DiffResults.gradient, result) == grads + + tape = ReverseDiff.GradientTape(g251, (x, y)) + for t in (tape, ReverseDiff.compile(tape)) + result = (DiffResults.GradientResult(MVector(x)), DiffResults.GradientResult(MVector(y))) + result = ReverseDiff.gradient!(result, t, (x, y)) + @test map(DiffResults.value, result) == (value, value) + @test map(DiffResults.gradient, result) == grads + end + + # a result tuple that does not match the input tuple is rejected by dispatch + @test_throws MethodError ReverseDiff.gradient!((MVector(x),), g251, (x, y)) + @test_throws MethodError ReverseDiff.gradient!((MVector(x), MVector(y), MVector(x)), g251, (x, y)) +end + +############################################################################################ + +# The output does not depend on the input, so it is recorded untracked and all derivatives +# are zero. +f_untracked(x) = 1.0 +g_untracked(x, y) = 2.0 + +@testset "output that does not depend on the input" begin + x, y = rand(3), rand(2) + + @test ReverseDiff.gradient(f_untracked, x) == zeros(3) + @test ReverseDiff.gradient(g_untracked, (x, y)) == (zeros(3), zeros(2)) + + result = ReverseDiff.gradient!(DiffResults.GradientResult(x), f_untracked, x) + @test DiffResults.value(result) == 1.0 + @test DiffResults.gradient(result) == zeros(3) + + result = (DiffResults.GradientResult(x), DiffResults.GradientResult(y)) + result = ReverseDiff.gradient!(result, g_untracked, (x, y)) + @test map(DiffResults.value, result) == (2.0, 2.0) + @test map(DiffResults.gradient, result) == (zeros(3), zeros(2)) + + tape = ReverseDiff.GradientTape(g_untracked, (x, y)) + for t in (tape, ReverseDiff.compile(tape)) + @test ReverseDiff.gradient!(t, (x, y)) == (zeros(3), zeros(2)) + @test ReverseDiff.gradient!((similar(x), similar(y)), t, (x, y)) == (zeros(3), zeros(2)) + end + + # a result that does not match the input tuple is rejected by dispatch, as when tracked + @test_throws MethodError ReverseDiff.gradient!((similar(x),), g_untracked, (x, y)) + @test_throws MethodError ReverseDiff.gradient!(similar(x), g_untracked, (x, y)) + @test_throws MethodError ReverseDiff.gradient!(similar(x), g251, (x, y)) + @test_throws MethodError ReverseDiff.gradient!(DiffResults.GradientResult(x), g251, (x, y)) +end + end # module diff --git a/test/api/HessianTests.jl b/test/api/HessianTests.jl index f489403e..c1c49a0a 100644 --- a/test/api/HessianTests.jl +++ b/test/api/HessianTests.jl @@ -1,6 +1,6 @@ module HessianTests -using DiffTests, ForwardDiff, ReverseDiff, Test +using DiffTests, ForwardDiff, ReverseDiff, StaticArrays, Test include(joinpath(dirname(@__FILE__), "../utils.jl")) @@ -106,4 +106,28 @@ end @test ReverseDiff.hessian(fa, x) == [12.0 0.0; 0.0 18.0] end +############################################################################################ + +f251(x) = sum(abs2, x) + +@testset "primal value of an immutable result (#251)" begin + x = MVector{2}(3.0, 5.0) + value, grad, hess = 34.0, [6.0, 10.0], [2.0 0.0; 0.0 2.0] + + result = DiffResults.HessianResult(MVector(x)) + result = ReverseDiff.hessian!(result, f251, x, ReverseDiff.HessianConfig(result, x)) + @test result isa DiffResults.ImmutableDiffResult + @test DiffResults.value(result) == value + @test DiffResults.gradient(result) == grad + @test DiffResults.hessian(result) == hess + + tape = ReverseDiff.HessianTape(f251, x) + for t in (tape, ReverseDiff.compile(tape)) + result = ReverseDiff.hessian!(DiffResults.HessianResult(MVector(x)), t, x) + @test DiffResults.value(result) == value + @test DiffResults.gradient(result) == grad + @test DiffResults.hessian(result) == hess + end +end + end # module diff --git a/test/api/JacobianTests.jl b/test/api/JacobianTests.jl index 507edc7b..0f26db34 100644 --- a/test/api/JacobianTests.jl +++ b/test/api/JacobianTests.jl @@ -1,6 +1,6 @@ module JacobianTests -using DiffTests, ForwardDiff, ReverseDiff, Test +using DiffTests, ForwardDiff, ReverseDiff, StaticArrays, Test include(joinpath(dirname(@__FILE__), "../utils.jl")) @@ -312,4 +312,25 @@ for f in DiffTests.BINARY_MATRIX_TO_MATRIX_FUNCS # test_approx(Jb test_b) end +############################################################################################ + +f251(x) = x .^ 2 + +@testset "primal value of an immutable result (#251)" begin + x = MVector{2}(3.0, 5.0) + value, jac = [9.0, 25.0], [6.0 0.0; 0.0 10.0] + + result = ReverseDiff.jacobian!(DiffResults.JacobianResult(SVector{2}(0.0, 0.0), x), f251, x) + @test result isa DiffResults.ImmutableDiffResult + @test DiffResults.value(result) == value + @test DiffResults.jacobian(result) == jac + + tape = ReverseDiff.JacobianTape(f251, x) + for t in (tape, ReverseDiff.compile(tape)) + result = ReverseDiff.jacobian!(DiffResults.JacobianResult(SVector{2}(0.0, 0.0), x), t, x) + @test DiffResults.value(result) == value + @test DiffResults.jacobian(result) == jac + end +end + end # module