diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1ad1ff41..90022731 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -48,6 +48,8 @@ jobs: - uses: julia-actions/cache@v3 - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 + env: + JULIA_NUM_THREADS: 4 - uses: julia-actions/julia-processcoverage@v1 - uses: codecov/codecov-action@v7 with: diff --git a/src/api/Config.jl b/src/api/Config.jl index d28198b9..afd2b261 100644 --- a/src/api/Config.jl +++ b/src/api/Config.jl @@ -105,7 +105,6 @@ function JacobianConfig(output::AbstractArray{D}, input::Tuple, tp::InstructionT return _JacobianConfig(cfg_input, cfg_output, tp) end -# we dispatch on V<:Real here because InstructionTape is actually also an AbstractArray function JacobianConfig(output::AbstractArray{D}, input::AbstractArray{V}, tp::InstructionTape = InstructionTape()) where {D, V <: Real} cfg_input = track(similar(input), D, tp) cfg_output = track!(similar(output, TrackedReal{D, D, Nothing}), output, tp) diff --git a/src/api/tape.jl b/src/api/tape.jl index d3e290b8..3442be37 100644 --- a/src/api/tape.jl +++ b/src/api/tape.jl @@ -101,8 +101,8 @@ methods on each instruction in the tape. function CompiledTape(t::T) where {T <: AbstractTape} return CompiledTape{T}( t, - [FunctionWrapper{Nothing, Tuple{}}(ForwardExecutor(instruction)) for instruction in t.tape], - [FunctionWrapper{Nothing, Tuple{}}(ReverseExecutor(t.tape[i])) for i in length(t.tape):-1:1] + [FunctionWrapper{Nothing, Tuple{}}(ForwardExecutor(instruction)) for instruction in instructions(t.tape)], + [FunctionWrapper{Nothing, Tuple{}}(ReverseExecutor(instruction)) for instruction in Iterators.reverse(instructions(t.tape))] ) end @@ -204,6 +204,7 @@ See `ReverseDiff.gradient` for a description of acceptable types for `input`. function GradientTape(f, input, cfg::GradientConfig = GradientConfig(input)) track!(cfg.input, input) tracked_ouput = f(cfg.input) + finish!(cfg.tape) return _GradientTape(f, cfg.input, tracked_ouput, cfg.tape) end @@ -212,6 +213,7 @@ function GradientTape(f, input::Tuple, cfg::GradientConfig = GradientConfig(inpu track!(cfg.input[i], input[i]) end tracked_output = f(cfg.input...) + finish!(cfg.tape) return _GradientTape(f, cfg.input, tracked_output, cfg.tape) end @@ -234,6 +236,7 @@ See `ReverseDiff.jacobian` for a description of acceptable types for `input`. function JacobianTape(f, input, cfg::JacobianConfig = JacobianConfig(input)) track!(cfg.input, input) tracked_ouput = f(cfg.input) + finish!(cfg.tape) return _JacobianTape(f, cfg.input, tracked_ouput, cfg.tape) end @@ -242,6 +245,7 @@ function JacobianTape(f, input::Tuple, cfg::JacobianConfig = JacobianConfig(inpu track!(cfg.input[i], input[i]) end tracked_output = f(cfg.input...) + finish!(cfg.tape) return _JacobianTape(f, cfg.input, tracked_output, cfg.tape) end @@ -261,6 +265,7 @@ function JacobianTape(f!, output, input, cfg::JacobianConfig = JacobianConfig(ou track!(cfg.output, output, cfg.tape) track!(cfg.input, input) f!(cfg.output, cfg.input) + finish!(cfg.tape) return _JacobianTape(f!, cfg.input, cfg.output, cfg.tape) end @@ -270,6 +275,7 @@ function JacobianTape(f!, output, input::Tuple, cfg::JacobianConfig = JacobianCo track!(cfg.input[i], input[i]) end f!(cfg.output, cfg.input...) + finish!(cfg.tape) return _JacobianTape(f!, cfg.input, cfg.output, cfg.tape) end @@ -296,5 +302,6 @@ function HessianTape(f, input, cfg::HessianConfig = HessianConfig(input)) track!(ht.input, input) gt = GradientTape(f, ht.input, gcfg) seeded_reverse_pass!(ht.output, gt.output, gt.input, gt.tape) + finish!(jcfg.tape) return ht end diff --git a/src/tape.jl b/src/tape.jl index be8e9323..9b27e78a 100644 --- a/src/tape.jl +++ b/src/tape.jl @@ -4,10 +4,75 @@ abstract type AbstractInstruction end -const InstructionTape = Vector{AbstractInstruction} +mutable struct TapeNode + const instruction::AbstractInstruction + prev::Union{Nothing, TapeNode} +end + +struct Finished end + +# A tape is first recorded, possibly from multiple threads at once, and then replayed. +# Reading it while recording throws, and so does recording onto it after `finish!`; +# `empty!` starts recording again. +mutable struct InstructionTape + # while recording: the most recently recorded node (nodes link backwards); + # `Finished()` once recording has finished + @atomic last::Union{Nothing, TapeNode, Finished} + # recorded instructions in order, filled by `finish!` + const instructions::Vector{AbstractInstruction} +end + +InstructionTape() = InstructionTape(nothing, AbstractInstruction[]) + +# end recording; all tasks recording onto `tp` must have finished (e.g. via `@sync`) +function finish!(tp::InstructionTape) + node = @atomicswap tp.last = Finished() + if node isa Finished + throw(ArgumentError("tape has already finished recording")) + end + while node !== nothing + push!(tp.instructions, node.instruction) + node = node.prev + end + reverse!(tp.instructions) + return tp +end + +# the recorded instructions; throws unless recording has finished +function instructions(tp::InstructionTape) + if !((@atomic tp.last) isa Finished) + throw(ArgumentError("tape is still recording; call `ReverseDiff.finish!` first")) + end + return tp.instructions +end + +Base.length(tp::InstructionTape) = length(instructions(tp)) + +function Base.empty!(tp::InstructionTape) + # `NULL_TAPE` never records, so it stays finished + if tp !== NULL_TAPE + empty!(tp.instructions) + @atomic tp.last = nothing + end + return tp +end + +# undo the swap in `record!` so that the tape stays finished +@noinline function throw_finished(tp::InstructionTape) + @atomic tp.last = Finished() + throw(ArgumentError("tape has finished recording; call `empty!` to record again")) +end -function record!(tp::InstructionTape, ::Type{InstructionType}, args...) where {InstructionType} - tp !== NULL_TAPE && push!(tp, InstructionType(args...)) +@inline function record!(tp::InstructionTape, ::Type{InstructionType}, args...) where {InstructionType <: AbstractInstruction} + if tp !== NULL_TAPE + node = TapeNode(InstructionType(args...), nothing) + # the swap orders concurrent recordings and returns the predecessor + prev = @atomicswap tp.last = node + if prev isa Finished + throw_finished(tp) + end + node.prev = prev + end return nothing end @@ -75,7 +140,7 @@ end ########## function forward_pass!(tape::InstructionTape) - for instruction in tape + for instruction in instructions(tape) forward_exec!(instruction) end return nothing @@ -85,8 +150,8 @@ end @noinline forward_exec!(instruction::SpecialInstruction) = special_forward_exec!(instruction) function reverse_pass!(tape::InstructionTape) - for i in length(tape):-1:1 - reverse_exec!(tape[i]) + for instruction in Iterators.reverse(instructions(tape)) + reverse_exec!(instruction) end return nothing end @@ -98,28 +163,41 @@ end # Pretty Printing # ################### -# extra spaces here accomodates padding in show(::IO, ::AbstractInstruction) -compactrepr(x::Tuple) = "(" * join(map(compactrepr, x), ",\n ") * ")" -compactrepr(x::AbstractArray) = length(x) < 5 ? match(r"\[.*?\]", repr(x)).match : summary(x) -compactrepr(x) = repr(x) - -function Base.show(io::IO, instruction::AbstractInstruction, pad = "") - name = isa(instruction, ScalarInstruction) ? "ScalarInstruction" : "SpecialInstruction" - println(io, pad, "$(name)($(instruction.func)):") - println(io, pad, " input: ", compactrepr(instruction.input)) - println(io, pad, " output: ", compactrepr(instruction.output)) - print(io, pad, " cache: ", compactrepr(instruction.cache)) +function Base.show(io::IO, instruction::AbstractInstruction) + print(io, nameof(typeof(instruction)), "(", instruction.func, ")") + return nothing +end + +function Base.show(io::IO, ::MIME"text/plain", instruction::AbstractInstruction) + show(io, instruction) + ctx = IOContext(io, :compact => true, :limit => true) + print(ctx, ":\n input: ") + show(ctx, instruction.input) + print(ctx, "\n output: ") + show(ctx, instruction.output) + print(ctx, "\n cache: ") + show(ctx, instruction.cache) return nothing end function Base.show(io::IO, tp::InstructionTape) - println(io, length(tp), "-element InstructionTape:") - i = 1 - for instruction in tp - print(io, "$i => ") - show(io, instruction) - println(io) - i += 1 + if (@atomic tp.last) isa Finished + print(io, length(tp), "-element InstructionTape") + else + print(io, "InstructionTape (recording)") + end + return nothing +end + +# like arrays, so that long tapes are truncated +function Base.show(io::IO, ::MIME"text/plain", tp::InstructionTape) + show(io, tp) + if (@atomic tp.last) isa Finished + instrs = instructions(tp) + if !isempty(instrs) + println(io, ":") + Base.print_array(io, instrs) + end end - return + return nothing end diff --git a/src/tracked.jl b/src/tracked.jl index 8f9c7f2f..d28d3c98 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -3,7 +3,7 @@ ######### const NULL_INDEX = typemin(Int) -const NULL_TAPE = InstructionTape() +const NULL_TAPE = finish!(InstructionTape()) # TrackedReal # #-------------# diff --git a/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index ffb6a9f9..7433d78b 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -134,6 +134,7 @@ end xt = ReverseDiff.track(copy(x), tp) # record yt = f(xt) + ReverseDiff.finish!(tp) @test length(tp) == 3 # sum, broadcast+, broadcast*, but not `g` end diff --git a/test/MacrosTests.jl b/test/MacrosTests.jl index a803f8cb..6e4ad358 100644 --- a/test/MacrosTests.jl +++ b/test/MacrosTests.jl @@ -41,18 +41,18 @@ function test_forward(f, x, tp) xt = ReverseDiff.TrackedReal(x, zero(x), tp) y = f(x) - @test isempty(tp) + @test isempty(take_recorded!(tp)) yt = f(xt) @test yt == y dual = f(Dual(x, one(x))) - @test length(tp) == 1 - instruction = first(tp) + recorded = take_recorded!(tp) + @test length(recorded) == 1 + instruction = first(recorded) @test typeof(instruction) <: ReverseDiff.ScalarInstruction @test instruction.input === xt @test instruction.output === yt @test instruction.cache[] === partials(dual, 1) - empty!(tp) return nothing end @@ -62,37 +62,37 @@ function test_forward(f, a, b, tp) c = f(a, b) dual = f(Dual(a, one(a), zero(a)), Dual(b, zero(b), one(b))) - @test isempty(tp) + @test isempty(take_recorded!(tp)) tc = f(at, b) @test tc == c - @test length(tp) == 1 - instruction = first(tp) + recorded = take_recorded!(tp) + @test length(recorded) == 1 + instruction = first(recorded) @test typeof(instruction) <: ReverseDiff.ScalarInstruction @test instruction.input === (at, b) @test instruction.output === tc @test instruction.cache[] === SVector(partials(dual, 1), partials(dual, 1)) - empty!(tp) tc = f(a, bt) @test tc == c - @test length(tp) == 1 - instruction = first(tp) + recorded = take_recorded!(tp) + @test length(recorded) == 1 + instruction = first(recorded) @test typeof(instruction) <: ReverseDiff.ScalarInstruction @test instruction.input === (a, bt) @test instruction.output === tc @test instruction.cache[] === SVector(partials(dual, 2), partials(dual, 2)) - empty!(tp) tc = f(at, bt) @test tc == c - @test length(tp) == 1 - instruction = first(tp) + recorded = take_recorded!(tp) + @test length(recorded) == 1 + instruction = first(recorded) @test typeof(instruction) <: ReverseDiff.ScalarInstruction @test instruction.input === (at, bt) @test instruction.output === tc @test instruction.cache[] === SVector(partials(dual)...) - empty!(tp) return nothing end @@ -142,11 +142,11 @@ function test_skip(g, x, tp) xt = ReverseDiff.TrackedReal(x, zero(x), tp) y = g(x) - @test isempty(tp) + @test isempty(take_recorded!(tp)) yt = g(xt) @test yt === y - @test isempty(tp) + @test isempty(take_recorded!(tp)) return nothing end @@ -155,19 +155,19 @@ function test_skip(g, a, b, tp) bt = ReverseDiff.TrackedReal(b, zero(b), tp) c = g(a, b) - @test isempty(tp) + @test isempty(take_recorded!(tp)) tc = g(at, b) @test tc === c - @test isempty(tp) + @test isempty(take_recorded!(tp)) tc = g(a, bt) @test tc === c - @test isempty(tp) + @test isempty(take_recorded!(tp)) tc = g(at, bt) @test tc === c - @test isempty(tp) + @test isempty(take_recorded!(tp)) return nothing end diff --git a/test/TapeTests.jl b/test/TapeTests.jl index 1a8f2ed0..d22d1fa5 100644 --- a/test/TapeTests.jl +++ b/test/TapeTests.jl @@ -22,19 +22,45 @@ for Instr in (SpecialInstruction, ScalarInstruction) tp = InstructionTape() ReverseDiff.record!(tp, Instr, +, (x, y, k), z, c) - @test tp[1] == instr - @test tp[1].func === + - @test tp[1].input[1] !== x - @test tp[1].input[2] !== y - @test tp[1].input[3] === k - @test tp[1].input[1] == x - @test tp[1].input[2] == y - @test tp[1].output === z - @test tp[1].cache === c - @test startswith(string(tp), "1-element InstructionTape:") + @test_throws ArgumentError length(tp) + @test_throws ArgumentError ReverseDiff.instructions(tp) + @test string(tp) == "InstructionTape (recording)" + @test repr(MIME"text/plain"(), tp) == "InstructionTape (recording)" + ReverseDiff.finish!(tp) + @test length(tp) == 1 + recorded = only(ReverseDiff.instructions(tp)) + @test recorded == instr + @test recorded.func === + + @test recorded.input[1] !== x + @test recorded.input[2] !== y + @test recorded.input[3] === k + @test recorded.input[1] == x + @test recorded.input[2] == y + @test recorded.output === z + @test recorded.cache === c + name = string(nameof(Instr), "(+)") + @test string(recorded) == name + @test startswith(repr(MIME"text/plain"(), recorded), name * ":\n input: ") + @test string(tp) == "1-element InstructionTape" + @test repr(MIME"text/plain"(), tp) == "1-element InstructionTape:\n " * name + @test_throws ArgumentError ReverseDiff.record!(tp, Instr, +, (x, y, k), z, c) + # the rejected instruction is not recorded and the tape stays finished + @test length(tp) == 1 + @test_throws ArgumentError ReverseDiff.finish!(tp) + + empty!(tp) + ReverseDiff.record!(tp, Instr, +, (x, y, k), z, c) + ReverseDiff.record!(tp, Instr, +, (x, y, k), z, c) + ReverseDiff.finish!(tp) + @test length(tp) == 2 ReverseDiff.record!(NULL_TAPE, Instr, +, (x, y, k), z, c) - @test isempty(NULL_TAPE) + @test length(NULL_TAPE) == 0 + @test repr(MIME"text/plain"(), NULL_TAPE) == "0-element InstructionTape" + # `NULL_TAPE` stays finished + empty!(NULL_TAPE) + @test length(NULL_TAPE) == 0 + @test repr(MIME"text/plain"(), NULL_TAPE) == "0-element InstructionTape" end end # module diff --git a/test/TrackedTests.jl b/test/TrackedTests.jl index b9c12ea9..19e90514 100644 --- a/test/TrackedTests.jl +++ b/test/TrackedTests.jl @@ -612,13 +612,13 @@ A = TrackedArray{BigInt, Float64, 1, Array{BigInt, 1}, Array{Float64, 1}} T = TrackedReal{BigInt, Float64, A} t2 = convert(TrackedReal{BigFloat, BigFloat, Nothing}, ta1) -@test length(tp) == 1 -instr = tp[1] +recorded = take_recorded!(tp) +@test length(recorded) == 1 +instr = first(recorded) @test instr.func === convert @test instr.input === ta1 @test samefields(instr.output, TrackedReal(big(varr[1]), big(darr[1]), tp)) @test instr.cache === nothing -empty!(tp) @test_throws ArgumentError convert(Float64, tr) @test_throws ArgumentError convert(BigFloat, tr) @@ -668,38 +668,38 @@ ta_sub = ta[:, :] idx = ReverseDiff.index_iterable(axes(ta), (:, :)) @test collect(idx) == [(i, j) for i in 1:3, j in 1:3] @test samefields(ta_sub, ta) -@test length(tp) == 1 -instr = tp[1] +recorded = take_recorded!(tp) +@test length(recorded) == 1 +instr = first(recorded) @test instr.func === getindex @test instr.input === (ta, idx) @test samefields(instr.output, ta) @test instr.cache === nothing -empty!(tp) for T in (UInt, Int) ta_sub = ta[:, T(1):T(2)] idx = ReverseDiff.index_iterable(axes(ta), (:, T(1):T(2))) @test collect(idx) == [(i, j) for i in 1:3, j in 1:2] @test samefields(ta_sub, TrackedArray(varr[:, 1:2], darr[:, 1:2], tp)) - @test length(tp) == 1 - instr = tp[1] + recorded = take_recorded!(tp) + @test length(recorded) == 1 + instr = first(recorded) @test instr.func === getindex @test instr.input === (ta, idx) @test samefields(instr.output, TrackedArray(varr[:, 1:2], darr[:, 1:2], tp)) @test instr.cache === nothing - empty!(tp) ta_sub = ta[T(2):T(3), :] idx = ReverseDiff.index_iterable(axes(ta), (T(2):T(3), :)) @test collect(idx) == [(i, j) for i in 2:3, j in 1:3] @test samefields(ta_sub, TrackedArray(varr[2:3, :], darr[2:3, :], tp)) - @test length(tp) == 1 - instr = tp[1] + recorded = take_recorded!(tp) + @test length(recorded) == 1 + instr = first(recorded) @test instr.func === getindex @test instr.input === (ta, idx) @test samefields(instr.output, TrackedArray(varr[2:3, :], darr[2:3, :], tp)) @test instr.cache === nothing - empty!(tp) S = T === UInt ? Int : UInt for U in (S, T) @@ -707,39 +707,39 @@ for T in (UInt, Int) idx = ReverseDiff.index_iterable(axes(ta), (S(1):S(2), T(2):T(3))) @test collect(idx) == [(i, j) for i in 1:2, j in 2:3] @test samefields(ta_sub, TrackedArray(varr[1:2, 2:3], darr[1:2, 2:3], tp)) - @test length(tp) == 1 - instr = tp[1] + recorded = take_recorded!(tp) + @test length(recorded) == 1 + instr = first(recorded) @test instr.func === getindex @test instr.input === (ta, idx) @test samefields(instr.output, TrackedArray(varr[1:2, 2:3], darr[1:2, 2:3], tp)) @test instr.cache === nothing - empty!(tp) end ta_sub = ta[T(2):T(6)] idx = ReverseDiff.index_iterable(axes(ta), (T(2):T(6),)) @test collect(idx) == [(i,) for i in 2:6] @test samefields(ta_sub, TrackedArray(varr[2:6], darr[2:6], tp)) - @test length(tp) == 1 - instr = tp[1] + recorded = take_recorded!(tp) + @test length(recorded) == 1 + instr = first(recorded) @test instr.func === getindex @test instr.input === (ta, idx) @test samefields(instr.output, TrackedArray(varr[2:6], darr[2:6], tp)) @test instr.cache === nothing - empty!(tp) end ta_sub = ta[:] idx = ReverseDiff.index_iterable(axes(ta), (:,)) @test collect(idx) == [(i, j) for i in 1:3, j in 1:3] @test samefields(ta_sub, TrackedArray(varr[:], darr[:], tp)) -@test length(tp) == 1 -instr = tp[1] +recorded = take_recorded!(tp) +@test length(recorded) == 1 +instr = first(recorded) @test instr.func === getindex @test instr.input === (ta, idx) @test samefields(instr.output, TrackedArray(varr[:], darr[:], tp)) @test instr.cache === nothing -empty!(tp) # logical indices (`Bool <: Integer`, so they reach the generic `getindex`) rowmask, colmask = [true, false, true], BitVector([false, true, true]) @@ -747,8 +747,7 @@ mask = [true false true; false true false; true false true] for inds in ((rowmask, :), (:, colmask), (rowmask, 2:3), (mask,), (vec(mask),)) @test samefields(@inferred(ta[inds...]), TrackedArray(varr[inds...], darr[inds...], tp)) - @test length(tp) == 1 - empty!(tp) + @test length(take_recorded!(tp)) == 1 end # a 0-d array has no index to dispatch on (#126) @@ -757,19 +756,18 @@ ta0 = TrackedArray(fill(varr[1]), fill(darr[1]), tp0) @test samefields(ta0[], ta0[1]) @test samefields(ta0[], ta0[CartesianIndex()]) @test ta0[].origin === ta0 -@test isempty(tp0) +@test isempty(take_recorded!(tp0)) # `view` aliases the parent's buffers, so nothing needs to be recorded ta_view = @inferred view(ta, :, 2) @test samefields(ta_view, TrackedArray(varr[:, 2], darr[:, 2], tp)) -@test isempty(tp) +@test isempty(take_recorded!(tp)) @test parent(ReverseDiff.value(ta_view)) === varr @test parent(ReverseDiff.deriv(ta_view)) === darr # views that aren't `IndexLinear` are materialized through `getindex` instead @test samefields(@inferred(view(ta, 1:2, :)), TrackedArray(varr[1:2, :], darr[1:2, :], tp)) -@test length(tp) == 1 -empty!(tp) +@test length(take_recorded!(tp)) == 1 # recording a view must not scale with its length function view_allocs(d) diff --git a/test/api/GradientTests.jl b/test/api/GradientTests.jl index 76d2b4d5..e101ab7f 100644 --- a/test/api/GradientTests.jl +++ b/test/api/GradientTests.jl @@ -490,4 +490,32 @@ end @test_throws ArgumentError ReverseDiff.gradient(W -> loss(zeros(3, 5), W, H), W) end +@testset "recording from multiple threads (#218)" begin + function g(x, i) + acc = zero(eltype(x)) + for _ in 1:100 + acc += (i * x[1] - x[2])^2 + end + return acc + sum(abs2, i .* x) + end + function f(x) + s = Vector{Any}(undef, 64) + @sync for i in eachindex(s) + Threads.@spawn s[i] = g(x, i) + end + return sum(s) + end + f_serial(x) = sum(i -> g(x, i), 1:64) + + x = [1.0, 2.0] + ∇f = ReverseDiff.gradient(f_serial, x) + for _ in 1:20 + @test ReverseDiff.gradient(f, x) ≈ ∇f + end + + # tape replay + tape = ReverseDiff.compile(ReverseDiff.GradientTape(f, x)) + @test ReverseDiff.gradient!(similar(x), tape, x) ≈ ∇f +end + end # module diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index cdaef6a4..7905e3a5 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -29,6 +29,7 @@ function test_elementwise(f, fopt, x, tp) # record yt = map(fopt, xt) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -52,6 +53,7 @@ function test_elementwise(f, fopt, x, tp) # record yt = broadcast(fopt, xt) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -79,6 +81,7 @@ function test_map(f, fopt, a, b, tp) # record ct = map(fopt, at, b) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -100,6 +103,7 @@ function test_map(f, fopt, a, b, tp) # record ct = map(fopt, a, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -121,6 +125,7 @@ function test_map(f, fopt, a, b, tp) # record ct = map(fopt, at, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -164,6 +169,7 @@ function test_broadcast(f, fopt, a::AbstractArray, b::AbstractArray, tp, builtin # record ct = g(at, b) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -185,6 +191,7 @@ function test_broadcast(f, fopt, a::AbstractArray, b::AbstractArray, tp, builtin # record ct = g(a, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -206,6 +213,7 @@ function test_broadcast(f, fopt, a::AbstractArray, b::AbstractArray, tp, builtin # record ct = g(at, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -250,6 +258,7 @@ function test_broadcast(f, fopt, n::Number, x::AbstractArray, tp, builtin::Bool # record yt = g(nt, x) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -271,6 +280,7 @@ function test_broadcast(f, fopt, n::Number, x::AbstractArray, tp, builtin::Bool # record yt = g(n, xt) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -292,6 +302,7 @@ function test_broadcast(f, fopt, n::Number, x::AbstractArray, tp, builtin::Bool # record yt = g(nt, xt) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -334,6 +345,7 @@ function test_broadcast(f, fopt, x::AbstractArray, n::Number, tp, builtin::Bool # record yt = g(x, nt) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -355,6 +367,7 @@ function test_broadcast(f, fopt, x::AbstractArray, n::Number, tp, builtin::Bool # record yt = g(xt, n) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -376,6 +389,7 @@ function test_broadcast(f, fopt, x::AbstractArray, n::Number, tp, builtin::Bool # record yt = g(xt, nt) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse diff --git a/test/derivatives/LinAlgTests.jl b/test/derivatives/LinAlgTests.jl index 6b961c58..6a627263 100644 --- a/test/derivatives/LinAlgTests.jl +++ b/test/derivatives/LinAlgTests.jl @@ -21,10 +21,11 @@ function test_arr2num(f, x, tp; ignore_tape_length = false) # record yt = f(xt) + ReverseDiff.finish!(tp) if !ignore_tape_length @test length(tp) == 1 end - @test abs(yt - y) <= abs(y) * eps(typeof(y)) + @test abs(value(yt) - y) <= abs(y) * eps(typeof(y)) # reverse ReverseDiff.seed!(yt) @@ -50,6 +51,7 @@ function test_arr2arr(f, x, tp) # record yt = f(xt) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -77,6 +79,7 @@ function test_arr2arr(f, a, b, tp) # record ct = f(at, b) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -98,6 +101,7 @@ function test_arr2arr(f, a, b, tp) # record ct = f(a, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -119,6 +123,7 @@ function test_arr2arr(f, a, b, tp) # record ct = f(at, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -152,6 +157,7 @@ function test_arr2arr_inplace(f!, f, c, a, b, tp) ct = track(c, eltype(c), ReverseDiff.NULL_TAPE) f!(ct, at, b) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -174,6 +180,7 @@ function test_arr2arr_inplace(f!, f, c, a, b, tp) ct = track(c, eltype(c), ReverseDiff.NULL_TAPE) f!(ct, a, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -196,6 +203,7 @@ function test_arr2arr_inplace(f!, f, c, a, b, tp) ct = track(c, eltype(c), ReverseDiff.NULL_TAPE) f!(ct, at, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse diff --git a/test/derivatives/ScalarTests.jl b/test/derivatives/ScalarTests.jl index 3bb2ae60..f1e35aa2 100644 --- a/test/derivatives/ScalarTests.jl +++ b/test/derivatives/ScalarTests.jl @@ -23,6 +23,7 @@ function test_forward(f, x, tp::InstructionTape, fsym::Symbol) # record yt = f(xt) @test yt == y + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -51,6 +52,7 @@ function test_forward(f, a, b, tp) # record ct = f(at, b) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -73,6 +75,7 @@ function test_forward(f, a, b, tp) # record ct = f(a, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -95,6 +98,7 @@ function test_forward(f, a, b, tp) # record ct = f(at, bt) @test ct == c + ReverseDiff.finish!(tp) @test length(tp) == 1 # reverse @@ -123,7 +127,7 @@ function test_skip(f, x, tp) y = f(x) yt = f(xt) @test yt == y - @test isempty(tp) + @test isempty(take_recorded!(tp)) return nothing end @@ -134,15 +138,15 @@ function test_skip(f, a, b, tp) ct = f(at, b) @test ct == c - @test isempty(tp) + @test isempty(take_recorded!(tp)) ct = f(a, bt) @test ct == c - @test isempty(tp) + @test isempty(take_recorded!(tp)) ct = f(at, bt) @test ct == c - @test isempty(tp) + @test isempty(take_recorded!(tp)) return nothing end diff --git a/test/utils.jl b/test/utils.jl index c21e491e..9bd8b052 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -17,6 +17,13 @@ test_println(kind, f, pad = " ") = println(pad, "testing $(kind): `$(f)`...") @inline test_approx(A, B, _atol = 1.0e-5; nans::Bool = false) = @test isapprox(A, B; atol = _atol, nans = nans) +# finish recording onto `tp`, return the recorded instructions, and start recording again +function take_recorded!(tp::InstructionTape) + instructions = copy(ReverseDiff.instructions(ReverseDiff.finish!(tp))) + empty!(tp) + return instructions +end + tracked_is(a, b) = value(a) === value(b) && deriv(a) === deriv(b) && tape(a) === tape(b) tracked_is(a::AbstractArray, b::AbstractArray) = all(map(tracked_is, a, b)) tracked_is(a::Tuple, b::Tuple) = all(map(tracked_is, a, b))