Why Turing.jl does not work with NUTS() and HMC() samplers?

General Usage
2

This is common issue I’ve stumbled on when using ForwardDiff.jl. From what I can tell, when a ForwardDiff.Dual is added/subtracted/multiplied with an indexed vector of Float64, there is an attempted conversion to allow for the operation (Conversion and Promotion · The Julia Language). ForwardDiff deals with this gracefully when you use R = alpha .+ t but as soon as indexing is involved, it struggles. Someone with a better understanding than me will have to explain why that is. (I’m guessing it has to do with making temporary copies of the arrays during the operation).

My usual course of action is to try to remove indexed arrays in my code or to ensure that they take on the ForwardDiff.Dual element type. So, for example, in your code t[tb:end] and R[tb:end] have elements of types Int64 and Float64 respectively. If you do the following:

using Turing

@model function ModelTuring(y, N::Int, t)

    α ~ Distributions.Uniform(0, 1)

    #########################################################
    # Ensure arrays have the proper type for ForwardDiff.Dual
    R = zeros(eltype(α), length(t))
    t = convert(Array{eltype(α)}, t)
    #########################################################

    tb=1
    ## This works fine with all samplers
    # R = α .+ t

    ## This only works with the MH() sampler only
    R[tb:end] = α .+ t[tb:end]

    for i=1:N
        y[i] ~ Poisson(R[i])
    end

end ##

you can ensure that the elements of the arrays are correctly typed. There may be a more elegant/performant way, but I haven’t found it yet.

3 Likes