tests: add tests for layers LinearReadout and ESNCell

Author: MartinuzziFrancesco

Project.toml

@@ -4,7 +4,6 @@ authors = ["Francesco Martinuzzi"]
 version = "0.11.4"
 
 [deps]
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
@@ -29,7 +28,6 @@ RCMLJLinearModelsExt = "MLJLinearModels"
 RCSparseArraysExt = "SparseArrays"
 
 [compat]
-Adapt = "4.1.1"
 Aqua = "0.8"
 ArrayInterface = "7.19.0"
 CellularAutomata = "0.0.6"
@@ -64,4 +62,4 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test", "SafeTestsets", "DifferentialEquations", "MLJLinearModels", "LIBSVM", "Statistics", "SparseArrays", "CellularAutomata", "JET"]
+test = ["Aqua", "Test", "SafeTestsets", "DifferentialEquations", "MLJLinearModels", "LIBSVM", "Statistics", "SparseArrays", "JET"]

docs/src/api/layers.md

@@ -3,9 +3,10 @@
 ## Base Layers
 ```@doc
     ReservoirChain
-    Readout
     Collect
     StatefulLayer
+    LinearReadout
+    SVMReadout
 ```
 
 ## Echo State Networks

docs/src/api/models.md

@@ -3,4 +3,5 @@
 ```@docs
     ESN
     DeepESN
+    HybridESN
 ```

src/ReservoirComputing.jl

@@ -1,6 +1,5 @@
 module ReservoirComputing
 
-using Adapt: adapt
 using ArrayInterface: ArrayInterface
 using Compat: @compat
 using ConcreteStructs: @concrete
@@ -15,13 +14,13 @@ using Static: StaticBool, StaticInt, StaticSymbol,
 using Reexport: Reexport, @reexport
 using WeightInitializers: DeviceAgnostic, PartialFunction, Utils
 @reexport using WeightInitializers
-@reexport using LuxCore: setup, apply
+@reexport using LuxCore: setup, apply, initialparameters, initialstates
 
 const BoolType = Union{StaticBool,Bool,Val{true},Val{false}}
 const InputType = Tuple{<:AbstractArray,Tuple{<:AbstractArray}}
 const IntegerType = Union{Integer,StaticInteger}
 
-#@compat(public, (create_states)) #do I need to add intialstates/parameters in compat?
+#@compat(public, (initialparameters)) #do I need to add intialstates/parameters in compat?
 
 #layers
 include("layers/basic.jl")
@@ -42,7 +41,7 @@ include("models/hybridesn.jl")
 #extensions
 include("extensions/reca.jl")
 
-export ESNCell, StatefulLayer, Readout, ReservoirChain, Collect, collectstates, train!, predict
+export ESNCell, StatefulLayer, LinearReadout, ReservoirChain, Collect, collectstates, train!, predict
 export SVMReadout
 export Pad, Extend, NLAT1, NLAT2, NLAT3, PartialSquare, ExtendedSquare
 export StandardRidge

src/inits/esn_inits.jl

@@ -787,7 +787,7 @@ function pseudo_svd(rng::AbstractRNG, ::Type{T}, dims::Integer...;
     reservoir_matrix = create_diag(rng, T, res_dim, T(max_value);
         sorted=sorted, reverse_sort=reverse_sort)
 
-    tmp = get_sparsity(R, res_dim)
+    tmp = get_sparsity(reservoir_matrix, res_dim)
     while tmp <= sparsity
         i = rand_range(rng, res_dim)
         j = rand_range(rng, res_dim)
@@ -813,7 +813,7 @@ end
 
 function create_diag(rng::AbstractRNG, ::Type{T}, res_dim::Integer, max_value::Number;
     sorted::Bool=true, reverse_sort::Bool=false) where {T<:Number}
-    diag_matrix = DeviceAgnostic.rand(rng, T, Int(n)) .* T(max_value)
+    diag_matrix = DeviceAgnostic.rand(rng, T, Int(res_dim)) .* T(max_value)
 
     if sorted
         sort!(diag_matrix)

src/layers/basic.jl

@@ -2,10 +2,10 @@ abstract type AbstractReservoirCollectionLayer <: AbstractLuxLayer end
 abstract type AbstractReservoirRecurrentCell <: AbstractLuxLayer end
 abstract type AbstractReservoirTrainableLayer <: AbstractLuxLayer end
 
-### Readout
+### LinearReadout
 # adapted from lux layers/basic Dense
 @doc raw"""
-    Readout(in_dims => out_dims, [activation];
+    LinearReadout(in_dims => out_dims, [activation];
             use_bias=false, include_collect=true)
 
 Linear readout layer with optional bias and elementwise activation. Intended as
@@ -48,7 +48,7 @@ before this layer (logically inserting a [`Collect()`](@ref) right before it).
   Otherwise training may operate on the post-readout signal,
   which is usually unintended.
 """
-@concrete struct Readout <: AbstractReservoirTrainableLayer
+@concrete struct LinearReadout <: AbstractReservoirTrainableLayer
     activation
     in_dims <: IntegerType
     out_dims <: IntegerType
@@ -58,32 +58,32 @@ before this layer (logically inserting a [`Collect()`](@ref) right before it).
     include_collect <: StaticBool
 end
 
-function Readout(mapping::Pair{<:IntegerType,<:IntegerType}, activation=identity; kwargs...)
-    return Readout(first(mapping), last(mapping), activation; kwargs...)
+function LinearReadout(mapping::Pair{<:IntegerType,<:IntegerType}, activation=identity; kwargs...)
+    return LinearReadout(first(mapping), last(mapping), activation; kwargs...)
 end
 
-function Readout(in_dims::IntegerType, out_dims::IntegerType, activation=identity;
+function LinearReadout(in_dims::IntegerType, out_dims::IntegerType, activation=identity;
     init_weight=rand32, init_bias=rand32, include_collect::BoolType=True(),
     use_bias::BoolType=False())
-    return Readout(activation, in_dims, out_dims, init_weight, init_bias, static(use_bias), static(include_collect))
+    return LinearReadout(activation, in_dims, out_dims, init_weight, init_bias, static(use_bias), static(include_collect))
 end
 
-function initialparameters(rng::AbstractRNG, ro::Readout)
+function initialparameters(rng::AbstractRNG, ro::LinearReadout)
     weight = ro.init_weight(rng, ro.out_dims, ro.in_dims)
 
     if has_bias(ro)
-        return (; weight, bias=ro.init_bias(rng, Float32, ro.out_dims))
+        return (; weight, bias=ro.init_bias(rng, ro.out_dims))
     else
         return (; weight)
     end
 end
 
-parameterlength(ro::Readout) = ro.out_dims * ro.in_dims + has_bias(ro) * ro.out_dims
-statelength(ro::Readout) = 0
+parameterlength(ro::LinearReadout) = ro.out_dims * ro.in_dims + has_bias(ro) * ro.out_dims
+statelength(ro::LinearReadout) = 0
 
-outputsize(ro::Readout, _, ::AbstractRNG) = (ro.out_dims,)
+outputsize(ro::LinearReadout, _, ::AbstractRNG) = (ro.out_dims,)
 
-function (ro::Readout)(inp::AbstractArray, ps, st::NamedTuple)
+function (ro::LinearReadout)(inp::AbstractArray, ps, st::NamedTuple)
     out_tmp = ps.weight * inp
     if has_bias(ro)
         out_tmp += ps.bias
@@ -92,8 +92,8 @@ function (ro::Readout)(inp::AbstractArray, ps, st::NamedTuple)
     return output, st
 end
 
-function Base.show(io::IO, ro::Readout)
-    print(io, "Readout($(ro.in_dims) => $(ro.out_dims)")
+function Base.show(io::IO, ro::LinearReadout)
+    print(io, "LinearReadout($(ro.in_dims) => $(ro.out_dims)")
     (ro.activation == identity) || print(io, ", $(ro.activation)")
     has_bias(ro) || print(io, ", use_bias=false")
     ic = known(getproperty(ro, Val(:include_collect)))
@@ -136,7 +136,7 @@ vectors are concatenated with `vcat` in order of appearance.
 
 ## Notes
 
-- When used with a single `Collect()` before a [`Readout`](@ref), training uses exactly
+- When used with a single `Collect()` before a [`LinearReadout`](@ref), training uses exactly
   the tensor right before the readout (e.g., the reservoir state).
 - With **multiple** `Collect()` layers (e.g., after different submodules), the
   per-step features are `vcat`-ed in chain order to form one feature vector.
@@ -150,7 +150,7 @@ vectors are concatenated with `vcat` in order of appearance.
           StatefulLayer(ESNCell(3 => 300)),
           NLAT2(),
           Collect(), # <-- collect the 300-dim reservoir after NLAT2
-          Readout(300 => 3; include_collect=false) # <-- toggle off the default Collect()
+          LinearReadout(300 => 3; include_collect=false) # <-- toggle off the default Collect()
       )
 ```
 """
@@ -173,7 +173,7 @@ in a step, the feature defaults to the final vector exiting the chain for
 that time step.
 
 !!! note
-    If your [`Readout`](@ref) layer was created with `include_collect=true`
+    If your [`LinearReadout`](@ref) layer was created with `include_collect=true`
     (default behaviour), a collection point is placed immediately before the readout,
     so the collected features are the inputs to the readout.
 
@@ -209,7 +209,12 @@ function collectstates(rc::AbstractLuxLayer, data::AbstractMatrix, ps, st::Named
         end
         push!(collected, state_vec === nothing ? copy(inp_tmp) : state_vec)
     end
-    states = eltype(data).(reduce(hcat, collected))
+    @assert !isempty(collected)
+    firstcol = collected[1]
+    Tcol = eltype(firstcol)
+    empty_mat = zeros(Tcol, length(firstcol), 0)
+    states_raw = reduce(hcat, collected; init=empty_mat)
+    states = eltype(data).(states_raw)
     return states, newst
 end
 

src/layers/lux_layers.jl

@@ -70,12 +70,12 @@ end
 
 wrap_functions_in_chain_call(x) = x
 
-_readout_include_collect(ro::Readout) = begin
+_readout_include_collect(ro::LinearReadout) = begin
     res = known(getproperty(ro, Val(:include_collect)))
     res === nothing ? false : res
 end
 
-function wrap_functions_in_chain_call(ro::Readout)
+function wrap_functions_in_chain_call(ro::LinearReadout)
     return _readout_include_collect(ro) ? (Collect(), ro) : ro
 end
 

src/models/deepesn.jl

@@ -53,7 +53,7 @@ function DeepESN(in_dims::Int,
         end
         prev = res_dims[res]
     end
-    ro = Readout(prev => out_dims, readout_activation)
+    ro = LinearReadout(prev => out_dims, readout_activation)
     return ReservoirChain((layers..., ro)...)
 end
 

src/models/esn.jl

@@ -5,7 +5,7 @@ function ESN(in_dims::IntegerType, res_dims::IntegerType, out_dims::IntegerType,
     cell = ESNCell(in_dims => res_dims, activation; kwargs...)
     mods = state_modifiers isa Tuple || state_modifiers isa AbstractVector ?
            Tuple(state_modifiers) : (state_modifiers,)
-    ro = Readout(res_dims => out_dims, readout_activation)
+    ro = LinearReadout(res_dims => out_dims, readout_activation)
     return ReservoirChain((StatefulLayer(cell), mods..., ro)...)
 end
 
@@ -18,7 +18,7 @@ end
 
 function Base.show(io::IO, ::MIME"text/plain", rc::ReservoirChain)
     L = collect(pairs(rc.layers))
-    if !isempty(L) && (L[1][2] isa StatefulLayer) && (L[end][2] isa Readout)
+    if !isempty(L) && (L[1][2] isa StatefulLayer) && (L[end][2] isa LinearReadout)
         sl = L[1][2]
         ro = L[end][2]
         if sl.cell isa ESNCell

src/models/hybridesn.jl

@@ -66,7 +66,7 @@ end
               kwargs...)
 
 Build a hybrid ESN as a `ReservoirChain`:
-`StatefulLayer(ESNCell) → modifiers → AttachStream(train KB) → Readout`.
+`StatefulLayer(ESNCell) → modifiers → AttachStream(train KB) → LinearReadout`.
 """
 function HybridESN(km::KnowledgeModel,
     in_dims::Integer, res_dims::Integer, out_dims::Integer,
@@ -82,7 +82,7 @@ function HybridESN(km::KnowledgeModel,
     stream_train = kb_stream_train(km, km.datasize)
     d_kb = size(stream_train, 1)
 
-    ro = Readout((res_dims + d_kb) => out_dims, readout_activation;
+    ro = LinearReadout((res_dims + d_kb) => out_dims, readout_activation;
         include_collect=static(include_collect))
 
     return ReservoirChain((StatefulLayer(cell), mods..., AttachStream(stream_train), ro)...)