forwarddiff.jl

module Loop
using ForwardDiff: Dual, Partials

function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
               seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
    for i in eachindex(duals)
        duals[i] = Dual{T,V,N}(x[i], seed)
    end
    return duals
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
               seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
    for i in 1:N
        duals[i] = Dual{T,V,N}(x[i], seeds[i])
    end
    return duals
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
               seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
    offset = index - 1
    for i in 1:N
        j = i + offset
        duals[j] = Dual{T,V,N}(x[j], seed)
    end
    return duals
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
               seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
    offset = index - 1
    for i in 1:chunksize
        j = i + offset
        duals[j] = Dual{T,V,N}(x[j], seeds[i])
    end
    return duals
end

end


module Broadcast
using ForwardDiff: Dual, Partials

function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
               seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
    duals .= Dual{T,V,N}.(x, Ref(seed))
    return duals
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
               seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
    duals[1:N] .= Dual{T,V,N}.(view(x, 1:N), seeds)
    return duals
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
               seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
    offset = index - 1
    chunk = (1:N) .+ offset
    duals[chunk] .= Dual{T,V,N}.(view(x, chunk), Ref(seed))
    return duals
end

function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
               seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
    offset = index - 1
    chunk = (1:chunksize) .+ offset
    if chunksize == N
        duals[chunk] .= Dual{T,V,N}.(view(x, chunk), seeds)
    else
        duals[chunk] .= Dual{T,V,N}.(view(x, chunk), getindex.(Ref(seeds), 1:chunksize))
    end
    return duals
end

end


using BenchmarkTools
using ForwardDiff: GradientConfig

function main()
    let
        result = nothing
        for M in (Loop, Broadcast)
            x = rand(1000)
            cfg = GradientConfig(nothing, x)

            println(@which M.seed!(cfg.duals, x, cfg.seeds[1]))

            result′ = M.seed!(cfg.duals, x, cfg.seeds[1])
            if result != nothing
                @test result == result′
            end

            @btime $M.seed!($(cfg.duals), $x, $(cfg.seeds[1]))
        end
    end

    println()

    let
        result = nothing
        for M in (Loop, Broadcast)
            x = rand(1000)
            cfg = GradientConfig(nothing, x)

            println(@which M.seed!(cfg.duals, x, cfg.seeds))

            result′ = M.seed!(cfg.duals, x, cfg.seeds)
            if result != nothing
                @test result == result′
            end

            @btime $M.seed!($(cfg.duals), $x, $(cfg.seeds))
        end
    end

    println()

    let
        result = nothing
        for M in (Loop, Broadcast)
            x = rand(1000)
            cfg = GradientConfig(nothing, x)

            println(@which M.seed!(cfg.duals, x, 2, cfg.seeds[1]))

            result′ = M.seed!(cfg.duals, x, 2, cfg.seeds[1])
            if result != nothing
                @test result == result′
            end

            @btime $M.seed!($(cfg.duals), $x, 2, $(cfg.seeds[1]))
        end
    end

    println()

    let
        result = nothing
        for M in (Loop, Broadcast)
            x = rand(1000)
            cfg = GradientConfig(nothing, x)

            println(@which M.seed!(cfg.duals, x, 2, cfg.seeds))

            result′ = M.seed!(cfg.duals, x, 2, cfg.seeds)
            if result != nothing
                @test result == result′
            end

            @btime $M.seed!($(cfg.duals), $x, 2, $(cfg.seeds))
        end
    end

    println()

    let
        result = nothing
        for M in (Loop, Broadcast)
            x = rand(1000)
            cfg = GradientConfig(nothing, x)

            println(@which M.seed!(cfg.duals, x, 2, cfg.seeds, 2))

            result′ = M.seed!(cfg.duals, x, 2, cfg.seeds, 2)
            if result != nothing
                @test result == result′
            end

            @btime $M.seed!($(cfg.duals), $x, 2, $(cfg.seeds), 2)
        end
    end

    return
end