Vector of Flux Model Parameters Sharing Underlying Data

Specific Domains Machine Learning
,
2

If you write append! this will produce a vector of all parameters, like destructure, but this vector won’t be tied to the original separate arrays – the storage for a Vector all needs to be in one place.

What you could do is the reverse, and re-make the model with views of one vector, which will still be tied:

function flat_and(model)
    arrays = AbstractVector[]
    fmap(model) do x
        x isa AbstractArray && push!(arrays, vec(x))
        x
    end
    flat = reduce(vcat, arrays)
    offset = Ref(0)
    out = fmap(model) do x
        x isa AbstractArray || return x
        y = view(flat, offset[] .+ (1:length(x)))
        offset[] += length(x)
        return reshape(y, size(x))
    end
    flat, out
end

m = Dense(2,3)
v, m2 = flat_and(m)
v[1] = 999
m2.weight[1,1] == 999  # true