using Flux
using Functors
using Optimisers

struct Custom
    abc::Tuple
end
Functors.@functor Custom (abc,)

function (f::Custom)(x)
    x .* f.abc[1] .+ f.abc[2]
end

function Custom(;dim::Int)
    abc = (randn(Float32, dim), randn(Float32, dim))
    return Custom(abc)
end

model = Flux.Chain(Dense(4, 16), Custom(;dim=16), Dense(16, 4))
p, re = Optimisers.destructure(model)
re(p)(randn(Float32, 4, 1))

ERROR: LoadError: type Tuple has no field layers
Stacktrace:
  [1] getproperty
    @ ./Base.jl:42 [inlined]
  [2] functor(#unused#::Type{Chain{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Custom, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}}, c::Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}, NamedTuple{(:abc,), Tuple{Tuple{Int64, Int64}}}, NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}})
    @ Flux ~/.julia/packages/Flux/qAdFM/src/layers/basic.jl:44
  [3] _trainable_biwalk(f::Function, x::Chain{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Custom, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, aux::Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}, NamedTuple{(:abc,), Tuple{Tuple{Int64, Int64}}}, NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}})
    @ Optimisers ~/.julia/packages/Optimisers/UAVzc/src/destructure.jl:94
  [4] #fmap#30
    @ ~/.julia/packages/Functors/qBIlC/src/functor.jl:78 [inlined]
  [5] _rebuild(x::Chain{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Custom, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, off::Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}, NamedTuple{(:abc,), Tuple{Tuple{Int64, Int64}}}, NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}}, flat::Vector{Float32}, len::Int64; walk::Function, kw::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Optimisers ~/.julia/packages/Optimisers/UAVzc/src/destructure.jl:83
  [6] _rebuild
    @ ~/.julia/packages/Optimisers/UAVzc/src/destructure.jl:82 [inlined]
  [7] (::Optimisers.Restructure{Chain{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Custom, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}, Tuple{NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}, NamedTuple{(:abc,), Tuple{Tuple{Int64, Int64}}}, NamedTuple{(:weight, :bias, :σ), Tuple{Int64, Int64, Tuple{}}}}})(flat::Vector{Float32})
    @ Optimisers ~/.julia/packages/Optimisers/UAVzc/src/destructure.jl:51