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

[MartinuzziFrancesco]

Since the design of Lux models has changed a bit in the meantime this new approach should work now

using LuxCore, Random, ConcreteStructs, Static, WeightInitializers, ArrayInterface
using Static: StaticBool, StaticInt, StaticSymbol,
    True, False, static, known, dynamic, StaticInteger

import LuxCore: AbstractLuxLayer, AbstractLuxContainerLayer, AbstractLuxWrapperLayer,
    initialparameters, initialstates, parameterlength, statelength, outputsize,
    setup, apply, replicate

const IntegerType = Union{Integer, StaticInteger}
const BoolType = Union{StaticBool, Bool, Val{true}, Val{false}}
const InputType = Tuple{<:AbstractMatrix, Tuple{<:AbstractMatrix}}

### from Lux.Utils
function sample_replicate(rng::AbstractRNG)
    rand(rng)
    return LuxCore.replicate(rng)
end

function init_hidden_state(rng::AbstractRNG, rnn, inp::AbstractMatrix)
    y = similar(inp, rnn.out_dims, Base.size(inp, 2))
    copyto!(y, rnn.init_state(rng, size(y)...))
    return ArrayInterface.aos_to_soa(y)
end
#from Lux extended_ops
const KnownSymbolType{v} = Union{Val{v}, StaticSymbol{v}}

function has_bias(l::AbstractLuxLayer)
    res = known(getproperty(l, Val(:use_bias)))
    return ifelse(res === nothing, false, res)
end

function getproperty(x, ::KnownSymbolType{v}) where {v}
    return v ∈ Base.propertynames(x) ? Base.getproperty(x, v) : nothing
end
###

@concrete struct ESNCell <: LuxCore.AbstractLuxLayer
    activation
    in_dims <: IntegerType
    out_dims <: IntegerType
    init_bias
    init_reservoir
    init_input
    #init_feedback::F
    init_state
    leak_coefficient
    use_bias <: StaticBool
end

function ESNCell((in_dims, out_dims)::Pair{<:Int, <:Int}, activation=tanh;
        use_bias::BoolType=True(), init_bias=zeros32, init_reservoir=glorot_uniform,
        init_input=glorot_uniform, init_state=zeros32, leak_coefficient = 1.0)
    return ESNCell(activation, in_dims, out_dims, init_bias, init_reservoir,
        init_input, init_state, leak_coefficient, use_bias)
end

function initialparameters(rng::AbstractRNG, esn::ESNCell)
    ps = (input_matrix = esn.init_input(rng, esn.out_dims, esn.in_dims),
        reservoir_matrix = esn.init_reservoir(rng, esn.out_dims, esn.out_dims),
        hidden_state = esn.init_state(rng, esn.out_dims))
    if has_bias(esn)
        ps = merge(ps, (bias = esn.init_bias(rng, esn.out_dims),))
    end
    return ps
end

function initialstates(rng::AbstractRNG, esn::ESNCell)
    return (rng = sample_replicate(rng),)
end

function (esn::ESNCell)(inp::AbstractMatrix, ps, st::NamedTuple)
    rng = replicate(st.rng)
    hidden_state = init_hidden_state(rng, esn, inp)
    return esn((inp, (hidden_state,)), ps, merge(st, (; rng)))
end

function (esn::ESNCell)((inp, (hidden_state,))::InputType, ps, st::NamedTuple)
    candidate_h = esn.activation.(ps.input_matrix * inp .+ 
        ps.reservoir_matrix * hidden_state .+ ps.bias)
    h_new = (1 - esn.leak_coefficient) .* hidden_state .+
        esn.leak_coefficient .* candidate_h
        
    return (h_new, (h_new,)), st
end