Replace ReverseDiff with Flux done (#457)

Author: xukai92

src/Turing.jl

@@ -20,15 +20,9 @@ using Markdown
   using Stan
   import Stan: Adapt, Hmc
 end
-@init @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
-  using ReverseDiff: GradientTape, GradientConfig, gradient!, compile, TrackedArray
-  import ReverseDiff: gradient
-end
-
 import Base: ~, convert, promote_rule, rand, getindex, setindex!
 import Distributions: sample
 import ForwardDiff: gradient
-import Flux: gradient
 using Flux: Tracker
 import MCMCChain: AbstractChains, Chains
 

src/core/ad.jl

@@ -108,8 +108,6 @@ gradient2(_vi::VarInfo, model::Function, spl::Union{Nothing, Sampler}) = begin
   g(vi[spl])
 end
 
-#  @init @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
-
 gradient_r(theta::Vector{Float64}, vi::VarInfo, model::Function) = gradient_r(theta, vi, model, nothing)
 gradient_r(theta::Vector{Float64}, vi::Turing.VarInfo, model::Function, spl::Union{Nothing, Sampler}) = begin
     f_r(ipts) = begin
@@ -119,12 +117,6 @@ gradient_r(theta::Vector{Float64}, vi::Turing.VarInfo, model::Function, spl::Uni
 
     grad = Tracker.gradient(f_r, theta)
 
-    # grad = ReverseDiff.gradient(x -> (vi[spl] = x; -runmodel(model, vi, spl).logp), inputs)
-
-    # vi[spl] = realpart(vi[spl])
-    # vi.logp = 0
-
-    map(x -> isa(x, Tracker.TrackedReal) ? x.data : x, grad)
+    first(grad).data
 end
 
-#  end

src/core/util.jl

@@ -4,10 +4,6 @@
 
 @inline invlogit(x::Union{T,Vector{T},Matrix{T}}) where T<:Real = one(T) ./ (one(T) .+ exp.(-x))
 @inline logit(x::Union{T,Vector{T},Matrix{T}}) where T<:Real = log.(x ./ (one(T) - x))
-@init @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
-  @inline invlogit(x::TrackedArray) = one(Real) ./ (one(Real) + exp.(-x))
-  @inline logit(x::TrackedArray) = log.(x ./ (one(Real) - x))
-end
 
 # More stable, faster version of rand(Categorical)
 function randcat(p::Vector{Float64})

src/helper.jl

@@ -10,9 +10,7 @@
 @inline realpart(ds::Matrix{Any}) = [realpart(col) for col in ds]
 @inline realpart(ds::Array)  = map(d -> realpart(d), ds)  # NOTE: this function is not optimized
 # @inline realpart(ds::TArray) = realpart(Array(ds))    # TODO: is it disabled temporary
-@init @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
-  @inline realpart(ta::ReverseDiff.TrackedReal) = ta.value
-end
+@inline realpart(ta::Tracker.TrackedReal) = ta.data
 
 @inline dualpart(d::ForwardDiff.Dual)       = d.partials.values
 @inline dualpart(ds::Union{Array,SubArray}) = map(d -> dualpart(d), ds)

src/samplers/hmcda.jl

@@ -112,8 +112,6 @@ function step(model, spl::Sampler{HMCDA}, vi::VarInfo, is_first::Bool)
       push!(spl.info[:accept_his], false)
 
       # Reset Θ
-      # NOTE: ForwardDiff and ReverseDiff need different implementation
-      #       due to struct Dual vs mutable TrackedReal
       if ADBACKEND == :forward_diff
 
         vi[spl] = θ
@@ -122,11 +120,7 @@ function step(model, spl::Sampler{HMCDA}, vi::VarInfo, is_first::Bool)
 
         vi_spl = vi[spl]
         for i = 1:length(θ)
-          if isa(vi_spl[i], ReverseDiff.TrackedReal)
-            vi_spl[i].value = θ[i]
-          else
-            vi_spl[i] = θ[i]
-          end
+          vi_spl[i] = θ[i]
         end
 
       end

src/samplers/support/hmc_core.jl

@@ -57,11 +57,7 @@ function gen_rev_func(vi, spl)
     elseif ADBACKEND == :reverse_diff
       vi_spl = vi[spl]
       for i = 1:length(θ_old)
-        if isa(vi_spl[i], ReverseDiff.TrackedReal)
-          vi_spl[i].value = θ_old[i]
-        else
-          vi_spl[i] = θ_old[i]
-        end
+        vi_spl[i] = θ_old[i]
       end
     end
     setlogp!(vi, old_logp)