Vector of Flux Model Parameters Sharing Underlying Data

I am trying to extract the parameters of a Flux model in a single vector such that they share the same underlying data, i.e. making changes to the vector changes the model parameters. The following is based on the Flux.destructure function:

function extract_params(m)
    xs = Float32[]
    fmap(m) do x
        x isa AbstractArray && push!(xs, vec(x))
        return x
    end
    return xs
end

However, this does not seem to work. I am somewhat new to Julia, so it is possible that I am misunderstanding some key concept, but it seems like this should be a relatively simple task. Can anyone help?

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

Thanks for the response! I did not realize that the storage for Vector needs to be contiguous, but I suppose that makes sense. This is a good workaround.