arraydist is showing unintended behavior.

[krishvishal]

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

[devmotion]

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.

[krishvishal]

Any examples for adjoints of Array{Type}? I don't know how to do that.

[mohamed82008]

This looks like an issue in DynamicPPL. u should be a TrackedArray here.

[mohamed82008]

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.

[mohamed82008]

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, ---)

[mohamed82008]

This TuringLang/Bijectors.jl#142 should fix this issue.