`unflatten` and mutation

[mhauru]

unflatten does not mutate its arguments. However, the result of unflatten shares memory with the argument.

vi = VarInfo()
vi = DynamicPPL.push!!(vi, @varname(a[1]), 1.0, Normal())
vi2 = DynamicPPL.unflatten(vi, [2.0])
vi2 = DynamicPPL.push!!(vi2, @varname(a[2]), 3.0, Normal())
@show keys(vi)
vi[@varname(a[2])]

prints keys(vi) = VarName[a[1], a[2]] but then crashes with a BoundsError on the last line.

I don't know if this is how most people use the !! convention, but I would argue that this would warrant calling unflatten unflatten!!. You have to treat it as if it may have mutated its arguments, because even if it hasn't yet, it's "tainted" the argument by linking it to the return value in a way that means you can't safely do anything with it anymore. Calling it unflatten!! would also give us permission to avoid an unnecessary call to copy in VNV, since we would explicitly allow mutating the input.

However, acknowledging that unflatten has this sort of behaviour would mess with how we use it in LogDensityFunction. The current model is that an LDF carries with it a persistent varinfo, and at every call to logdensity with new values x we do

    varinfo_new = unflatten(varinfo, x)
    varinfo_eval = last(evaluate!!(model, varinfo_new))
    return getlogdensity(varinfo_eval)

To be principled about this, we should instead make a copy of varinfo in each such call, or change unflatten so that it internally makes copies of the various data structures, in which case we wouldn't need to rename it. However, copies are not cheap, and it feels very unfortunate to have to do that given that we've been getting away with the current half-copying implementation of unflatten for so long: The actual values vector in varinfo and varinfo_new are distinct, only the surrounding structures like the dictionary of varnames are shared.

The right solution isn't immediately clear to me. I'll give this a think, thoughts welcome.