Tests using Flux.Tracker

[trappmartin]

We should add test which test agains the use of Flux.Tracker for AD. While playing around with the transformations I encountered some issues while using Flux.Tracker and simplex distributions.

[trappmartin]

Running the following code...

function f_test(θᴿ, x)
    K = 3
    μ = [0.5, 0.2, -0.5]
    @assert length(θᴿ) == K
    
    logpdf_mm(θ, x) = log(sum(θ .* (x.-μ).^2))
    
    d = Dirichlet(K, 1.0)
    lp = logpdf_with_trans(d, invlink(d, θᴿ), true)
    lp += logpdf_mm(invlink(d, θᴿ), x)
    return lp
end

v, g = Tracker.forward((θ, x) -> f_test(θ, x), [0.3, 0.2, 0.5], 0.4)

will result in:

MethodError: no method matching Float64(::Flux.Tracker.TrackedReal{Float64})
Closest candidates are:
  Float64(::Real, !Matched::RoundingMode) where T<:AbstractFloat at rounding.jl:185
  Float64(::T<:Number) where T<:Number at boot.jl:725
  Float64(!Matched::Int8) at float.jl:60
  ...

[trappmartin]

This issue is related to the fact that θ is TrackedArray, which for some reasons breaks the differentiation.

Stacktrace:
 [1] convert(::Type{Float64}, ::Flux.Tracker.TrackedReal{Float64}) at ./number.jl:7
 [2] setindex!(::Array{Float64,1}, ::Flux.Tracker.TrackedReal{Float64}, ::Int64) at ./array.jl:769
 [3] invlink(::Dirichlet{Float64}, ::TrackedArray{…,Array{Float64,1}}, ::Type{Val{true}}) at /Users/martin/.julia/dev/Bijectors/src/Bijectors.jl:228

[trappmartin]

If instead an Array of TrackedReal's is used, everything works smooth.

[mohamed82008]

This is fixed now.