DistributionsAD.jl
DistributionsAD.jl copied to clipboard
`arraydist` is showing unintended behavior.
using Turing, ReverseDiff, LoopVectorization
using ReverseDiff: @grad, TrackedArray
Turing.setadbackend(:reversediff)
logsumexp(x::TrackedArray) = ReverseDiff.track(logsumexp, x)
@grad function logsumexp(x::AbstractArray)
lse = logsumexp(ReverseDiff.value(x))
return lse, Δ -> (Δ .* exp.(ReverseDiff.value(x) .- lse),)
end
function logsumexp(mat)
maxima = vec(vreduce(max, mat, dims=1))'
exp_mat = @avx exp.(mat .- maxima)
sum_exp_mat = vreduce(+, exp_mat, dims=1)
return @avx sum_exp_mat .= maxima .+ log.(sum_exp_mat)
end
@model function mwe(y, A, ::Type{T} = Vector{Float64}) where {T}
n = size(A, 1)
scale ~ filldist(Uniform(1, 3), n)
TT = eltype(scale)
x ~ arraydist(truncated.(Laplace.(0, scale), TT(-10.0), TT(70.0)))
println(typeof(x))
μ = logsumexp(A .- x)[1, :]
y .~ Normal.(μ, 1)
return x
end
y1 = rand(10);
A1 = rand(10,10);
model = mwe(y1, A1);
chain = sample(model, NUTS(0.65), 100)
Running the above code results in the following output/error.
Terminal output:
Array{Float64,1}
Array{ReverseDiff.TrackedReal{Float64,Float64,Nothing},1}
ERROR: MethodError: no method matching typemin(::Type{ReverseDiff.TrackedReal{Float64,Float64,Nothing}})
Closest candidates are:
typemin(::Type{Bool}) at bool.jl:6
typemin(::Type{Int8}) at int.jl:651
typemin(::Type{UInt8}) at int.jl:653
...
Full stacktrace: https://pastebin.com/hWvFG3sG
The custom adjoint I've defined for logsumexp function expects TrackedArray type but x's type is Array{ReverseDiff.TrackedReal{Float64,Float64,Nothing},1}.
@mohamed82008
It seems the error arises from the use of LoopVectorization in your implementation of logsumexp which tries to call typemin which is not defined by ReverseDiff for ::Type{<:TrackedReal}. I guess you could either specify the custom adjoint for Array{<:TrackedReal} or not use LoopVectorization.
Any examples for adjoints of Array{Type}? I don't know how to do that.
This looks like an issue in DynamicPPL. u should be a TrackedArray here.
Ok so I narrowed down the problem here to a Bijectors-ReverseDiff issue. The problem is that for the arraydist of truncated Laplace above, the bijector is a TruncatedBijector with array lower and upper bounds. Applying this bijector to a TrackedArray gives an array of TrackedReal. I think this may be fixable.
julia> using Turing, ReverseDiff
julia> dist = arraydist(truncated.(Laplace.(0, [1, 2]), -10.0, 70.0));
julia> x = ReverseDiff.track(rand(dist));
julia> bijector(dist)(x)
2-element Array{ReverseDiff.TrackedReal{Float64,Float64,Nothing},1}:
TrackedReal<Fre>(-1.7994654328949322, 0.0, d6k, ---)
TrackedReal<Ggm>(-1.5788290245320968, 0.0, d6k, ---)
This https://github.com/TuringLang/Bijectors.jl/pull/142 should fix this issue.