Keep tracked values on the tape in deepcopy

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}

src/tape.jl

@@ -38,16 +38,27 @@ 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
 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

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)))

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,

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})

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