fix: small fixes and docs

Author: MartinuzziFrancesco

src/inits/esn_inits.jl

@@ -791,7 +791,7 @@ function pseudo_svd(rng::AbstractRNG, ::Type{T}, dims::Integer...;
     while tmp <= sparsity
         i = rand_range(rng, res_dim)
         j = rand_range(rng, res_dim)
-        θ = DeviceAgnostic.rand(rng, T) * T(2) - T(1)
+        θ = DeviceAgnostic.rand(rng, T) * T(2) .- T(1)
         reservoir_matrix = reservoir_matrix * create_qmatrix(rng, T, res_dim, i, j, θ)
         tmp = get_sparsity(reservoir_matrix, res_dim)
     end

src/layers/esn_cell.jl

@@ -1,3 +1,5 @@
+abstract type AbstractEchoStateNetworkCell <: AbstractReservoirRecurrentCell end
+
 @doc raw"""
     ESNCell(in_dims => out_dims, [activation];
         use_bias=false, init_bias=rand32,
@@ -63,7 +65,7 @@ Created by `initialstates(rng, esn)`:
 
   - `rng`: a replicated RNG used to sample initial hidden states when needed.
 """
-@concrete struct ESNCell <: AbstractReservoirRecurrentCell
+@concrete struct ESNCell <: AbstractEchoStateNetworkCell
     activation::Any
     in_dims <: IntegerType
     out_dims <: IntegerType
@@ -76,15 +78,15 @@ Created by `initialstates(rng, esn)`:
     use_bias <: StaticBool
 end
 
-function ESNCell(
-        (in_dims, out_dims)::Pair{<:IntegerType, <:IntegerType}, activation = tanh;
-        use_bias::BoolType = False(), init_bias = zeros32, init_reservoir = rand_sparse,
-        init_input = scaled_rand, init_state = randn32, leak_coefficient = 1.0)
+function ESNCell((in_dims, out_dims)::Pair{<:IntegerType, <:IntegerType},
+        activation = tanh; use_bias::BoolType = False(), init_bias = zeros32,
+        init_reservoir = rand_sparse, init_input = scaled_rand,
+        init_state = randn32, 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)
+function initialparameters(rng::AbstractRNG, esn::AbstractEchoStateNetworkCell)
     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))
     if has_bias(esn)
@@ -93,11 +95,11 @@ function initialparameters(rng::AbstractRNG, esn::ESNCell)
     return ps
 end
 
-function initialstates(rng::AbstractRNG, esn::ESNCell)
+function initialstates(rng::AbstractRNG, esn::AbstractEchoStateNetworkCell)
     return (rng = sample_replicate(rng),)
 end
 
-function (esn::ESNCell)(inp::AbstractArray, ps, st::NamedTuple)
+function (esn::AbstractEchoStateNetworkCell)(inp::AbstractArray, ps, st::NamedTuple)
     rng = replicate(st.rng)
     hidden_state = init_hidden_state(rng, esn, inp)
     return esn((inp, (hidden_state,)), ps, merge(st, (; rng)))

src/reservoircomputer.jl

@@ -1,3 +1,34 @@
+
+@doc raw"""
+    ReservoirComputer(reservoir, states_modifiers, readout)
+
+Generic reservoir-computing container that wires together:
+  1) a `reservoir` (any Lux-compatible layer producing features),
+  2) zero or more `states_modifiers` applied sequentially to the reservoir features,
+  3) a `readout` layer (typically [`LinearReadout`](@ref)).
+
+The container exposes a standard `(x, ps, st) -> (y, st′)` interface and
+utility functions to initialize parameters/states, stream sequences to collect
+features, and install trained readout weights.
+
+## Arguments
+
+- `reservoir`: a layer that consumes inputs and produces feature vectors.
+- `states_modifiers`: a tuple (or vector converted to `Tuple`) of layers applied
+  after the reservoir (may be empty).
+- `readout`: the final trainable layer mapping features to outputs.
+
+## Inputs
+
+- `x`: input to the reservoir (shape determined by the reservoir).
+- `ps`: reservoir computing parameters.
+- `st`: reservoir computing states.
+
+## Returns
+
+- `(y, st′)` where `y` is the readout output and `st′` contains the updated
+  states of the reservoir, modifiers, and readout.
+"""
 @concrete struct ReservoirComputer <: AbstractReservoirComputer{(:reservoir, :states_modifiers, :readout)}
     reservoir::Any
     states_modifiers::Any

test/esn/test_esn.jl

@@ -34,18 +34,18 @@ end
 
     ps, st = setup(rng, esn)
 
-    @test haskey(ps, :cell)
-    @test haskey(ps.cell, :input_matrix)
-    @test haskey(ps.cell, :reservoir_matrix)
-    @test !haskey(ps.cell, :bias)
-    @test size(ps.cell.input_matrix) == (res_dims, in_dims)
-    @test size(ps.cell.reservoir_matrix) == (res_dims, res_dims)
+    @test haskey(ps, :reservoir)
+    @test haskey(ps.reservoir, :input_matrix)
+    @test haskey(ps.reservoir, :reservoir_matrix)
+    @test !haskey(ps.reservoir, :bias)
+    @test size(ps.reservoir.input_matrix) == (res_dims, in_dims)
+    @test size(ps.reservoir.reservoir_matrix) == (res_dims, res_dims)
 
     @test haskey(ps, :readout)
     @test haskey(ps.readout, :weight)
     @test size(ps.readout.weight) == (out_dims, res_dims)
 
-    @test haskey(st, :cell)
+    @test haskey(st, :reservoir)
     @test haskey(st, :states_modifiers)
     @test haskey(st, :readout)
     @test st.states_modifiers isa Tuple
@@ -72,7 +72,7 @@ end
 
     @test size(Y) == (D, 1)
     @test vec(Y) ≈ x
-    @test haskey(st2, :cell) && haskey(st2, :states_modifiers) && haskey(st2, :readout)
+    @test haskey(st2, :reservoir) && haskey(st2, :states_modifiers) && haskey(st2, :readout)
 end
 
 @testset "ESN: forward (batch matrix) with identity pipeline -> Y == X" begin
@@ -155,7 +155,7 @@ end
     y, st2 = esn(x, ps, st)
 
     @test y ≈ x
-    @test haskey(st2, :cell)
+    @test haskey(st2, :reservoir)
     @test haskey(st2, :states_modifiers)
     @test haskey(st2, :readout)
 end

test/runtests.jl

@@ -3,7 +3,7 @@ using Test
 
 @testset "Common Utilities" begin
     @safetestset "Quality Assurance" include("qa.jl")
-    #@safetestset "States" include("test_states.jl")
+    @safetestset "States" include("test_states.jl")
 end
 
 @testset "Layers" begin