Build on LuxCore

[MartinuzziFrancesco]

A lot of recent issues can be tackled by an internal refactor of the drivers and generally of how the models are defined. Since this refactor is in the plans, it is worth exploring the possibility of building on LuxCore.jl, to provide a familiar interface to SciML users. I will update this issue with ideas on how to build the components

[MartinuzziFrancesco]

Some very raw initial ideas for the implementation:

struct ESNCell{use_bias, A, B, W, I, S} <: LuxCore.AbstractExplicitLayer
    activation::A
    in_dims::Int
    res_dims::Int
    init_bias::B
    init_reservoir::W
    init_input::I
    #init_feedback::F
    init_state::S
    leak_coefficient::Number
end

function ESNCell((in_dims, res_dims)::Pair{<:Int, <:Int}, activation=tanh;
    use_bias::Bool=false, init_bias=zeros32, init_reservoir=glorot_uniform,
    init_input=glorot_uniform, init_state=zeros32)
    return ESNCell{use_bias,typeof(activation),typeof(init_bias),
        typeof(init_reservoir),typeof(init_input),typeof(init_state)}(
        activation,in_dims,res_dims,init_bias,init_reservoir,init_input,init_state)
end

# none of the weights in a esn cell are trainable. should the be parameters or states?
function LuxCore.initialparameters(rng::AbstractRNG,
    esn::ESNCell{use_bias}) where {use_bias}

    ps = (input_matrix=esn.init_input(rng, esn.res_dims, esn.in_dims),
        reservoir_matrix=esn.init_reservoir(rng, esn.res_dims, esn.res_dims))
    use_bias && (ps = merge(ps, (bias=esn.init_bias(rng, esn.res_dims),)))
    return ps
end

function LuxCore.initialstates(rng::AbstractRNG,
    esn::ESNCell{use_bias}) where {use_bias}
    randn(rng, 1)
    st = (states = nothing, rng=LuxCore.replicate(rng))
    return st
end

const _ESNCellInputType = Tuple{<:AbstractMatrix, Tuple{<:AbstractMatrix}}

function (esn::ESNCell{true})((x, (hidden_state,))::_ESNCellInputType,
    ps, st::NamedTuple)
    h_new =(1 - esn.leaky_coefficient) * x + esn.activation.(st.input_matrix * x .+ 
        st.reservoir_matrix * hidden_state .+ st.bias)
    return (h_new, (h_new,)), st
end

function (esn::ESNCell{false})((x, (hidden_state,))::_ESNCellInputType,
    ps, st::NamedTuple)
    h_new = (1 - esn.leaky_coefficient) * x + esn.activation.(ps.input_matrix * x .+
        ps.reservoir_matrix * hidden_state)
    return (h_new, (h_new,)), st
end

# very WIP stuff here
# passes the data through the esn and obtains the states
function instantiate(esn::ESNCell, xs, ps, st)
    st_new = (:states)
    return NamedTuple(:last_state, :states)
end

struct ESN{T} <: LuxCore.AbstractExplicitContainerLayer{(:reservoir, :states, :readout)}
    reservoir::LuxCore.AbstractExplicitLayer
    states::Function # here go the nlas and state augmentations # this should be something like a chain
    readout::T # this has to be generic to accommodate multiple types of readouts # how do these return ps and st?
end

function ESN(in_size::Int, res_size::Int, out_size::Int;
    leaky_coefficient::Number = 0.0, init_reservoir = rand_sparse, init_input = scaled_rand,
    bias=false, init_bias=rand, readout = StandardRidge(0.0), states = States(), activation=tanh)

    esn = ESNCell(in_size, res_size, activation; bias=bias, init_reservoir=init_reservoir)
end

[MartinuzziFrancesco]

Actually I think all the weights in the ESN have to be states instead of parameters, in case someone needs to build a larger model with ESNs

function LuxCore.initialparameters(rng::AbstractRNG,
    esn::ESNCell{use_bias}) where {use_bias}

    return NamedTuple()
end

function LuxCore.initialstates(rng::AbstractRNG,
    esn::ESNCell{use_bias}) where {use_bias}
    randn(rng, 1)
    st = (input_matrix=esn.init_input(rng, esn.res_dims, esn.in_dims),
        reservoir_matrix=esn.init_reservoir(rng, esn.res_dims, esn.res_dims),
        states = nothing, rng=LuxCore.replicate(rng))
    use_bias && (st = merge(st, (bias=esn.init_bias(rng, esn.res_dims),)))
    return st
end