Variable ordering in VarInfo

[mhauru]

The order of the variables in a VarInfo is important, because in sampling we frequently vectorise and unflatten: vi_new = unflatten(vi, do_a_thing_with_vec(vi[:])). There's an implicit assumption in sample that the VarInfo has the same ordering from one iteration to the next. However, we don't have a very robust way of specifying the order in VarInfo. It's currently based on the order in which values are inserted into the VarInfo. subset used to mess with the order before #832, which caused a bug in Gibbs. subset isn't the only problem though, for instance the following could cause all sorts of trouble:

@model function f()
    z ~ Normal()
    if z > 0
        x ~ Normal(1.0)
        y ~ Normal(2.0)
    else
        y ~ Normal(2.0)
        x ~ Normal(1.0)
    end
end

In general the current reliance on insertion order feels fragile.

Could we develop a more robust approach to variable ordering? For instance, could we say that variables are ordered alphabetically by their VarName? This would make inserting new variables slower, but shouldn't have a performance impact once all variables are in place, which is the vast majority of sampling.

---

Only partially related, but note that even fixing the ordering wouldn't fix the below problem:

@model function f()
    z ~ Normal()
    if z > 0
        x_dim = 1
        y_dim = 2
    else
        x_dim = 2
        y_dim = 1
    end
    x = Vector{Float64}(undef, x_dim)
    y = Vector{Float64}(undef, y_dim)
    x .~ Normal(1.0)
    y .~ Normal(2.0)
end

Sampling this currently crashes for various reasons, but there's still a philosophical issue of if a sampler gives you a state for this model that is [0.1, 0.2, 0.3, 0.4], you can not know whether that should be taken as z = 0.1, x=[0.2], y=[0.3, 0.4] or z = 0.1, x=[0.2, 0.3], y=[0.4] without inspecting the f in detail. I'm not sure there's any way to fix this as long as samplers only see vectors of floats and are blind to varnames, but we should keep this case in mind in case we can think of a redesign that would fix it. Any such redesign would fix many variable ordering problems too.

[torfjelde]

I guess there are two aspects to this:

1. Does any other ordering fix the issue of deletion and insertion causing changes to the order?
2. Which ordering is the most intuitive?

I don't see how (1) would work, e.g. alphabetic ordering wouldn't address this 😕

For (2), I'm very pro the current approach: when I do varinfo[:][1] I'm expecting part of the first variable. IMO that ordering is the most intuitive.

Sampling this currently crashes for various reasons, but there's still a philosophical issue of if a sampler gives you a state for this model that is [0.1, 0.2, 0.3, 0.4], you can not know whether that should be taken as z = 0.1, x=[0.2], y=[0.3, 0.4] or z = 0.1, x=[0.2, 0.3], y=[0.4] without inspecting the f in detail

But this doesn't seem to be an ordering issue, no? Isn't this a huge part of the motivation for the introduction of VarNamedVector?