diff --git a/src/api/tape.jl b/src/api/tape.jl index 3442be37..5639b1bd 100644 --- a/src/api/tape.jl +++ b/src/api/tape.jl @@ -108,6 +108,15 @@ end Base.show(io::IO, t::CompiledTape) = print(io, typeof(t).name, "($(t.tape.func))") +# copied `FunctionWrapper`s keep pointers to the original instructions, so recompile instead +function Base.deepcopy_internal(t::CompiledTape, dict::IdDict) + return ( + get!(dict, t) do + CompiledTape(Base.deepcopy_internal(t.tape, dict)) + end + )::typeof(t) +end + const CompiledGradient{T <: GradientTape} = CompiledTape{T} const CompiledJacobian{T <: JacobianTape} = CompiledTape{T} const CompiledHessian{T <: HessianTape} = CompiledTape{T} diff --git a/src/tape.jl b/src/tape.jl index 9b27e78a..41a815ec 100644 --- a/src/tape.jl +++ b/src/tape.jl @@ -38,9 +38,11 @@ function finish!(tp::InstructionTape) return tp end +isrecording(tp::InstructionTape) = !((@atomic tp.last) isa Finished) + # the recorded instructions; throws unless recording has finished function instructions(tp::InstructionTape) - if !((@atomic tp.last) isa Finished) + if isrecording(tp) throw(ArgumentError("tape is still recording; call `ReverseDiff.finish!` first")) end return tp.instructions @@ -48,6 +50,15 @@ end Base.length(tp::InstructionTape) = length(instructions(tp)) +# `NULL_TAPE` is recognised by identity, so copies must keep it +function Base.deepcopy_internal(tp::InstructionTape, dict::IdDict) + if tp === NULL_TAPE + return tp + else + return invoke(Base.deepcopy_internal, Tuple{Any, IdDict}, tp, dict) + end +end + function Base.empty!(tp::InstructionTape) # `NULL_TAPE` never records, so it stays finished if tp !== NULL_TAPE diff --git a/src/tracked.jl b/src/tracked.jl index 7b592398..716b12f7 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -435,6 +435,16 @@ Base.size(t::TrackedArray) = size(value(t)) Base.copy(t::TrackedArray) = t +# while recording, tracked values are numbers on the tape: copying them would detach them; +# afterwards they are copied together with their tape +function Base.deepcopy_internal(t::Union{TrackedReal, TrackedArray}, dict::IdDict) + if isrecording(tape(t)) + return t + else + return invoke(Base.deepcopy_internal, Tuple{Any, IdDict}, t, dict) + end +end + Base.similar(t::TrackedArray, args::Union{Integer, AbstractUnitRange}...) = similar(value(t), eltype(t), args...) Base.similar(t::TrackedArray, T::Type, args::Union{Integer, AbstractUnitRange}...) = similar(value(t), T, args...) @@ -453,7 +463,6 @@ reshape_body = :(TrackedArray(reshape(value(t), dims), reshape(deriv(t), dims), Base.hash(t::TrackedReal, hsh::UInt64) = hash(value(t), hsh) -Base.deepcopy(t::T) where {T <: TrackedReal} = t Base.copy(t::T) where {T <: TrackedReal} = t Base.float(t::TrackedReal) = _float(t, float(value(t))) diff --git a/test/QATests.jl b/test/QATests.jl index 952e1d05..8da244cb 100644 --- a/test/QATests.jl +++ b/test/QATests.jl @@ -33,8 +33,8 @@ end UPSTREAM_INTERNALS = ( # `Base` and `Base.Broadcast` :AbstractArrayStyle, :Broadcasted, :DefaultArrayStyle, :LogicalIndex, :RefValue, - :combine_eltypes, :flatten, :instantiate, :print_array, :promote_op, :result_style, - :rtoldefault, + :combine_eltypes, :deepcopy_internal, :flatten, :instantiate, :print_array, :promote_op, + :result_style, :rtoldefault, # `DiffResults` :DiffResult, :GradientResult, :ImmutableDiffResult, :derivative, :gradient, :gradient!, :hessian, :jacobian, diff --git a/test/TrackedTests.jl b/test/TrackedTests.jl index 97be9f9f..0c535366 100644 --- a/test/TrackedTests.jl +++ b/test/TrackedTests.jl @@ -67,6 +67,8 @@ t = TrackedReal(v, d) @test t.value === v @test t.deriv === d @test t.tape === ReverseDiff.NULL_TAPE +# copies keep `NULL_TAPE` +@test deepcopy(t).tape === ReverseDiff.NULL_TAPE @test t.index === ReverseDiff.NULL_INDEX @test t.origin === nothing @test isa(t, TrackedReal{Float64, Int, Nothing}) diff --git a/test/api/GradientTests.jl b/test/api/GradientTests.jl index 8aa67c70..598e8d3f 100644 --- a/test/api/GradientTests.jl +++ b/test/api/GradientTests.jl @@ -487,3 +487,40 @@ end # a `Float64` buffer throws instead of zeroing the gradient @test_throws ArgumentError ReverseDiff.gradient(W -> loss(zeros(3, 5), W, H), W) end + +############################################################################################ + +deepcopy_real(x) = deepcopy(x[1]) * deepcopy(x[2]) +deepcopy_vector(x) = prod(deepcopy([x[1], x[2]])) +deepcopy_tuple(x) = prod(deepcopy((x[1], x[2]))) +deepcopy_namedtuple(x) = (y = deepcopy((; a = x[1], b = x[2])); y.a * y.b) +deepcopy_array(x) = prod(deepcopy(x)) +deepcopy_aliased(x) = x[1] * x[2] + x[1] - deepcopy((x[1],))[1] + +@testset "`deepcopy` keeps the tape (#327)" begin + for f in ( + deepcopy_real, deepcopy_vector, deepcopy_tuple, deepcopy_namedtuple, deepcopy_array, + deepcopy_aliased, + ) + @test ReverseDiff.gradient(f, [3.0, 4.0]) == [4.0, 3.0] + tape = ReverseDiff.GradientTape(f, [3.0, 4.0]) + for t in (tape, ReverseDiff.compile(tape)) + @test ReverseDiff.gradient!(t, [5.0, 2.0]) == [2.0, 5.0] + end + end +end + +@testset "`deepcopy` of a tape is independent (#327)" begin + tape = ReverseDiff.GradientTape(x -> sum(abs2, x), [3.0, 4.0]) + for t in (tape, ReverseDiff.compile(tape)) + copied = deepcopy(t) + @test ReverseDiff.gradient!(copied, [1.0, 2.0]) == [2.0, 4.0] + @test ReverseDiff.gradient!(t, [5.0, 6.0]) == [10.0, 12.0] + @test ReverseDiff.gradient!(copied, [1.0, 2.0]) == [2.0, 4.0] + + # aliased tapes stay aliased + copied = deepcopy((t, t)) + @test copied[1] === copied[2] + @test ReverseDiff.gradient!(copied[1], [1.0, 2.0]) == [2.0, 4.0] + end +end