ReverseDiff Error with Vector{Real} Input in Turing Model

Hi everyone!

We stumbled upon an issue when fitting a Turing model using ReverseDiff if the inputs are Vector{Real}.

Minimal Working Example

using Turing, ReverseDiff, Distributions, LogExpFunctions

@model function some_model(x, y)

    parameter ~ Uniform(0, 30)

    for i in eachindex(y)
        probabilities = softmax(x[i] * parameter, dims=1)
        distribution = Categorical(probabilities)
        y[i] ~ distribution 
    end
end

### This will work
log_probabilities = [
[-4.226834039860986, -1.10380534825695, -1.3485838864717796],
[-2.7411713024220434, -2.3756404263051873, -0.9992390187038857],
[-2.4084973027747805, -1.1944158586117868, -1.3887071333579988],
[-3.3499522854558443, -0.9086409330664779, -2.604388503142469],
[-2.3371073076367623, -1.3779453250505502, -1.0720136146812946]]

### This will crash 
log_probabilities = [
Real[-4.226834039860986, -1.10380534825695, -1.3485838864717796],
Real[-2.7411713024220434, -2.3756404263051873, -0.9992390187038857],
Real[-2.4084973027747805, -1.1944158586117868, -1.3887071333579988],
Real[-3.3499522854558443, -0.9086409330664779, -2.604388503142469],
Real[-2.3371073076367623, -1.3779453250505502, -1.0720136146812946]]

outcomes = [2,3,2,2,3]

my_model = some_model(log_probabilities, outcomes)

results = sample(my_model, NUTS(-1, 0.65; adtype=AutoReverseDiff(true)), 100)

Error Message:

In the case where our log_probabilities are Vector{Real} this results in the following error:

ERROR: MethodError: no method matching increment_deriv!(::Float64, ::Float64)

Closest candidates are:
  increment_deriv!(::ReverseDiff.TrackedReal, ::Real)
   @ ReverseDiff .julia\packages\ReverseDiff\p1MzG\src\derivatives\propagation.jl:45
  increment_deriv!(::AbstractArray, ::Any)
   @ ReverseDiff .julia\packages\ReverseDiff\p1MzG\src\derivatives\propagation.jl:38
  increment_deriv!(::ReverseDiff.TrackedArray, ::Real, ::Any)
   @ ReverseDiff .julia\packages\ReverseDiff\p1MzG\src\derivatives\propagation.jl:34

Now, the obvious solution is to simply keep our inputs as Vector{Float64} . However, this becomes an issue if we want to fit more complex models where we need to keep this as Vector{Real} .

As far as my understanding goes, the problem arises from ReverseDiff incorrectly tracking these inputs when it shouldn't, and consequently the wrong method for increment_deriv!() being called.

Feel free to move this issue to ReverseDiff if you think this is more ReverseDiff package specific.

Many thanks for your help!

TuringLang / Turing.jl

ReverseDiff Error with Vector{Real} Input in Turing Model #2364

Minimal Working Example