Do we still need parametric types in model definition for autodiff?

[yebai]

We currently use the following DynamicPPL syntax to support AutoDiff.

@model function f(N, ::Type{T}=Vector{Float64}) where {T}
    a = T(undef, N)
    a .~ Normal()
    m = sum(a)
    o = T(undef, N)
    o .~ Normal(m, 1)
end

We still need the parametric type (..., ::Type{T}=Vector{Float64}) where {T} for operator-overloading approaches like ForwardDiff and ReverseDiff. However, this is no longer necessary for Mooncake, which is source-transformation-based.

There might be a way to work around such requirements for ForwardDiff and ReverseDiff by implementing a (simple) model (AST) expression transform similar to Mooncake's. Such an approach is a bit speculative at the moment, but it is worth considering to simplify the modelling syntax.

cc @willtebbutt

[penelopeysm]

I started investigating this in #1015. I'm pretty optimistic that we don't need the type parameters, although there are some caveats. Here's a (long...) summary of my findings. I assume I should present this some day.

Tilde statements are OK

First of all, this model, which takes no special precautions with vector element types:

using Turing, ForwardDiff, ReverseDiff, Random, Accessors

@model function f1(y)
    x = Vector{Float64}(undef, length(y))
    for i in eachindex(y)
        x[i] ~ Normal()
        y[i] ~ Normal(x[i])
    end
end

actually runs perfectly fine on the current release of Turing (v0.39.10):

julia> mean(sample(Xoshiro(468), f1([1.0, 2.0]), NUTS(; adtype=AutoForwardDiff()), 1000; progress=false))
┌ Info: Found initial step size
└   ϵ = 1.6
Mean
  parameters      mean
      Symbol   Float64

        x[1]    0.4850
        x[2]    0.9886


julia> mean(sample(Xoshiro(468), f1([1.0, 2.0]), NUTS(; adtype=AutoReverseDiff()), 1000; progress=false))
┌ Info: Found initial step size
└   ϵ = 1.6
Mean
  parameters      mean
      Symbol   Float64

        x[1]    0.4850
        x[2]    0.9886

The key is that x[i] can only appear on the LHS of a tilde statement. This works because we use Accessors.set and that is able to set a ForwardDiff.Dual or ReverseDiff.TrackedReal into a Vector{Float64}, changing the type of the vector in the process.

julia> x = Float64[1.0, 2.0]
2-element Vector{Float64}:
 1.0
 2.0

julia> x = Accessors.set(x, (@optic _[1]), ForwardDiff.Dual(3.0)) # x[1] ~ dist actually does this!
2-element Vector{ForwardDiff.Dual{Nothing, Float64, 0}}:
 Dual{Nothing}(3.0)
 Dual{Nothing}(2.0)

Assignments are not OK

However, a model like this

@model function f2(y)
    x = Vector{Float64}(undef, length(y))
    s ~ Normal()
    for i in eachindex(y)
        x[i] = 2 * s
        y[i] ~ Normal(x[i])
    end
end

will fail with both ForwardDiff and ReverseDiff. The reason is because here x[i] is on the LHS of a plain =, and setindex! doesn't have the nice functionality of Accessors.set.

julia> x = Float64[1.0, 2.0]
2-element Vector{Float64}:
 1.0
 2.0

julia> x[1] = ForwardDiff.Dual(3.0)
ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{Nothing, Float64, 0})
The type `Float64` exists, but no method is defined for this combination of argument types when trying to construct it.

How to make assignments work

Now, there are perfectly good ways to work around this IMO. I've used Greek letters to number because, why not, I'm bored.

(α) Don't assign

The most basic option is to tell users to never assign into x[i], and to instead always use Accessors.set and/or promote_type. I'm actually not super opposed to this, but I'm going to list a bunch of alternatives in the hope of having a more meaningful discussion.

(β) Abstract types

This works:

@model function f3(y)
    x = Vector{Real}(undef, length(y))
    s ~ Normal()
    for i in eachindex(y)
        x[i] = 2 * s
        y[i] ~ Normal(x[i])
    end
end

This samples fine with both ForwardDiff and ReverseDiff but indexing into x is of course not type-stable.

(γ) Getting the type from a previous parameter

An improvement could be to use the type of s:

@model function f4(y)
    s ~ Normal()
    # Note that we had to move the definition of `x` to come after `s`.
    x = Vector{typeof(s)}(undef, length(y))
    for i in eachindex(y)
        x[i] = 2 * s
        y[i] ~ Normal(x[i])
    end
end

which also samples with ForwardDiff and ReverseDiff fine. If s is a vector or matrix one can just use eltype(s) instead.

IMO this workaround and in particular f4 is a very natural way to solve the problem. By natural I mean: it doesn't require internal DynamicPPL magic, and thus it is much more closely aligned with general Julia practices, and also much easier to explain to users.

Now, there're two problems. Firstly, although f4 samples fine, it gives different results with ForwardDiff and ReverseDiff. I have no idea why :(

julia> mean(sample(Xoshiro(468), f4([1.0, 2.0]), NUTS(; adtype=AutoForwardDiff()), 1000; progress=false))
┌ Info: Found initial step size
└   ϵ = 0.8
Mean
  parameters      mean
      Symbol   Float64

           s    0.6435


julia> mean(sample(Xoshiro(468), f4([1.0, 2.0]), NUTS(; adtype=AutoReverseDiff()), 1000; progress=false))
┌ Info: Found initial step size
└   ϵ = 0.8
Mean
  parameters      mean
      Symbol   Float64

           s    0.6511

Secondly, depending on how complicated the model is, it may require some non-trivial reorganisation. Consider this:

@model function f5(y)
    # v can just have Float64 because it is only on the LHS of tilde, not of assignment
    v = Vector{Float64}(undef, length(y))
    # But x has to be initialised before the loop
    x = Vector{???}(undef, length(y))
    for i in eachindex(y)
        # And v is only sampled in the loop so we can't access its type above
        v[i] ~ Normal()
        x[i] = 2 * v[i]
        y[i] ~ Normal(x[i])
    end
end

where ??? means I don't know how to get the right type.

This can be rewritten as:

@model function f6(y)
    v = Vector{Float64}(undef, length(y))
    for i in eachindex(y)
        v[i] ~ Normal()
    end
    # Now we can get the right type
    x = Vector{eltype(v)}(undef, length(y))
    for i in eachindex(y)
        x[i] = 2 * v[i]
        y[i] ~ Normal(x[i])
    end
end

This samples fine (and this time ForwardDiff and ReverseDiff give identical results; although it's different from f5 because of the reordering).

I don't know if this is sufficiently general. My impression is that all models can always be rewritten in this manner. Specifically, if x needs to be an AD tracer type, then it will depend on some random variable v, and the model can always be rearranged such that v is sampled before x is initialised. Conversely, if x does not depend on any random variable, then it doesn't need to be an AD tracer type, it can just be plain old Vector{Float64}.

I would love if people tried to come up with counterexamples to this!

(δ) eltype(__varinfo__)

To avoid having to rewrite the model, there is one remaining possibility: we can peer into the varinfo and check what type its values have. If it's being evaluated in the context of AD, then __varinfo__[:] will be some kind of AD tracer type from which we can infer the correct type to use. That is in fact almost exactly what we're doing right now, it's just that the ugliness is hidden behind the ::Type{T} stuff and matchingvalue.

So for example:

@model function f7(y)
    v = Vector{Float64}(undef, length(y))
    x = Vector{eltype(__varinfo__[:])}(undef, length(y))
    for i in eachindex(y)
        v[i] ~ Normal()
        x[i] = 2 * v[i]
        y[i] ~ Normal(x[i])
    end
end

I think perhaps this is where the idea of source transformation comes in. So we could, for example, write

@model function f7_macro(y)
    v = Vector{Float64}(undef, length(y))
    x <- newarray(length(y))
    for i in eachindex(y)
        v[i] ~ Normal()
        x[i] = 2 * v[i]
        y[i] ~ Normal(x[i])
    end
end

and have the model macro expand it to the above. Obviously, syntax can be bikeshedded. Alternatively, we could just make a macro specifically for this purpose.

@model function f7_macro_alt(y)
    v = Vector{Float64}(undef, length(y))
    x = @newarray(length(y))
    for i in eachindex(y)
        v[i] ~ Normal()
        x[i] = 2 * v[i]
        y[i] ~ Normal(x[i])
    end
end

[mhauru]

Not saying I'm necessarily supporting it as the choice here, but just noting that @newarrayor other new syntax for array creation would also help solve something in #965. Namely, I've been hoping we could say that x ~ dist is an assume statement if and only if x as a variable isn't set to have a value at that point in the model, otherwise it's an observe statement. This would go nicely with @yebai's latest proposal where parameter names are declared in a | (; a, b, c) block and not as the first argument of the model.

So

x = 1.0
x ~ Normal()  # observe

and

x ~ Normal()  # assume
y = 2*x
y ~ Normal()  # observe

I like the simplicity of this, because it makes implementation easy, but more importantly because it makes reasoning about models easy, while allowing very flexible model syntax.

The only problem I can see with that is array variables, which need to be initialised before individual elements can get their x[i] ~ dist statements. If we were to introduce a @newarray macro/new syntax for the purpose of AD type parameters, then we could also use it to deal with this problem, and have very simple semantics for what is an observation and what is an assumption.