Hi,
Destructuring struct with destructure
is very handy. Thank you.
However I’m a bit surprise that the restructuring is not Type stable. Am I missing something? Does it have some impact on speed when re(flat)
(and its gradient) is called very often ?
julia> using Optimisers
julia> struct Bar; a::Float64;b::Vector{Float64};end
julia> b = Bar(2.0, randn(10))
Bar(2.0, [0.42917489814746823, 1.266826972405323, 0.6992954064945902, -0.965529236441083, -0.6069994307534123, -0.5668205590293701, -0.231523643263828, 1.9408140168456018, 0.6864706994403353, 0.6929670854487652])
julia> flat, re = destructure(b)
([0.42917489814746823, 1.266826972405323, 0.6992954064945902, -0.965529236441083, -0.6069994307534123, -0.5668205590293701, -0.231523643263828, 1.9408140168456018, 0.6864706994403353, 0.6929670854487652], Restructure(Bar, ..., 10))
julia> @code_warntype re(flat)
MethodInstance for (::Optimisers.Restructure{Bar, @NamedTuple{a::Tuple{}, b::Int64}})(::Vector{Float64})
from (re::Optimisers.Restructure)(flat::AbstractVector) @ Optimisers ~/.julia/packages/Optimisers/W5seC/src/destructure.jl:59
Arguments
re::Optimisers.Restructure{Bar, @NamedTuple{a::Tuple{}, b::Int64}}
flat::Vector{Float64}
Body::Any
1 ─ %1 = Optimisers._rebuild::Core.Const(Optimisers._rebuild)
│ %2 = Base.getproperty(re, :model)::Bar
│ %3 = Base.getproperty(re, :offsets)::@NamedTuple{a::Tuple{}, b::Int64}
│ %4 = Base.getproperty(re, :length)::Int64
│ %5 = (%1)(%2, %3, flat, %4)::Any
└── return %5
Glad it’s useful. I agree it would be nice if this were type-stable.
I think the core problem is that Optimisers.jl cares about object identity. If the same mutable object appears more than once in a struct, this ===
relation is preserved. But this property isn’t visible to the type system. The reconstruction chooses whether to make a new Vector{Float64}
, or to retrieve one from an IdDict{Any}
, based on objectid
. (This is all outsourced to Functors.jl.)
That may not be insurmountable – the two choices should have the same type. There are examples of ID-aware recursion which are type-stable, e.g. Base.deepcopy
. In fact Functors.fmap
has been upgraded to often be type-stable. Maybe the same could be done for destructure
?
julia> struct Tri; a::Float64; b::Vector{Float64}; c::Vector{Float64}; end
julia> bi = Tri(1.11, [2.22, 3.33], [2.22, 3.33]);
julia> un = let twice = [2.22, 3.33]
Tri(1.11, twice, twice)
end;
julia> typeof(bi) == typeof(un)
true
julia> v1, re1 = Optimisers.destructure(un)
([2.22, 3.33], Restructure(Tri, ..., 2))
julia> v2, re2 = Optimisers.destructure(bi)
([2.22, 3.33, 2.22, 3.33], Restructure(Tri, ..., 4))
julia> typeof(re2) == typeof(re1)
true
julia> @code_warntype re1([4.4, 5.5]) # bad
julia> @code_warntype Base.deepcopy(un) # fine
julia> @code_warntype Functors.fmapstructure(complex, un) # also fine
See also issue 177.
Another approach is to make destructure!
which mutates the original object, instead of reconstructing. Then the return type ought to be trivial to infer. This was PR 165, it got stuck but could perhaps be revived.
Thank you very much. I didn’t saw this issue 177 that similar to my issue. For the moment, I will explicitly set the return type of the reconstruct expression as suggested there. I’m not guru enough to submit a PR on that.
Anyway, I find this package very useful and it has a broader use case than just optimization (in my case I don’t use the optimization part). May I suggest to put the parameter part in a separate package with a much clearer name for users. It will be easier to find.
Thanks again