Model no longer type stable when using destructure and restructure

[JTaets]

Hi,

When restructuring a model after Flux.destructure, it becomes type unstable (and slow).
Is there a way to avoid this?

using Flux
using BenchmarkTools

model = Chain(Dense(1,64,σ),Dense(64,1))
p, re = Flux.destructure(model)
re_model(x,p,re) = re(p)(x)  # passing re here doesn't make a difference

@code_warntype model([1f0])          # type stable
@btime $model($[1f0])                # 460.406 ns (4 allocations: 864 bytes)

@code_warntype re_model([1f0],p,re)  # type unstable
@btime $re_model($[1f0],$p,$re)      # 3.850 μs (57 allocations: 3.86 KiB)

[JTaets]

Like 1 second after posting i realize that re(p) is the thing that makes it slower

@btime $re($p) #3.237 μs (53 allocations: 3.02 KiB)

So i guess i'll restructure outside of re_model

[DhairyaLGandhi]

You should only need to restructure the model once and it would work like new.