no method matching ChainRules.OneElement

[paul188]

I am trying to code a function that casts a Vector{Float64} of weights and biases to ComponentArray that Lux.jl can handle. I have to do this in the optimization routine. Since this method involves mutating of arrays, I wrote my own rrule to handle it. Unfortunately, I encounter the error:

ERROR: LoadError: MethodError: no method matching ChainRules.OneElement(::Tangent{Any, NamedTuple{(:layer_1, :layer_2, :layer_3, :layer_4), NTuple{4, Tangent{Any, NamedTuple{(:weight, :bias), Tuple{Matrix{Float64}, Vector{Float64}}}}}}}, ::Tuple{Int64}, ::Tuple{Base.OneTo{Int64}})

Closest candidates are:
  ChainRules.OneElement(::T, ::I, ::A) where {N, T<:Number, I<:Tuple{Vararg{Int64, N}}, A<:Tuple{Vararg{AbstractUnitRange, N}}}

So it seems like the first argument ::T is missing
And I have tried quite long to fix it, but cannot find any mistakes in the code.
The conversion method from Vector to ComponentArray just distributes the vector entries into a nested structure, so it is basically the identity when viewed as a pullback. This is the method:

function convert_params_to_tuple_no_ode(p::Vector{Float64}, n_in::Int64, n_out::Int64, hidden_layers::Tuple{Int64,Int64})
    nn_subtuple = NamedTuple()
  ...
    return ComponentArrays.ComponentVector((ps_lux = nn_subtuple,))
end

And this is my attempt at the Pullback:

function ChainRules.rrule(::typeof(convert_params_to_tuple_no_ode),p::Vector{Float64}, n_in::Int64, n_out::Int64, hidden_layers::Tuple{Int64, Int64})
    Y = convert_params_to_tuple_no_ode(p, n_in, n_out, hidden_layers)
    Y2 = ComponentArrays._namedtuple(Y)

    function convert_params_to_tuple_no_ode_pullback(tuple_tangent)

        tuple_diff = Tangent{typeof(Y2)}(ps_lux = tuple_tangent.ps_lux)
        return (NoTangent(),tuple_diff, NoTangent(),NoTangent(),NoTangent())
    end
    return Y2, convert_params_to_tuple_no_ode_pullback
end

I converted everything to NamedTuples because the package seems unable to handle the ComponentArrays directly.

Would be very thankful for any help ! :) I hope this is the right outlet for my question.

[oxinabox]

I think probably better venue for your question is Discourse.
https://discourse.julialang.org/
More people will see it there.
Its a better place for help.