Type stability of destructure / re-build in Optimisers.jl

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