tests: add JET, small fixes

Author: MartinuzziFrancesco

Project.toml

@@ -36,6 +36,7 @@ CellularAutomata = "0.0.6"
 Compat = "4.16.0"
 ConcreteStructs = "0.2.3"
 DifferentialEquations = "7.16.1"
+JET = "0.9.20"
 LIBSVM = "0.8"
 LinearAlgebra = "1.10"
 LuxCore = "1.3.0"
@@ -54,6 +55,7 @@ julia = "1.10"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LIBSVM = "b1bec4e5-fd48-53fe-b0cb-9723c09d164b"
 MLJLinearModels = "6ee0df7b-362f-4a72-a706-9e79364fb692"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -62,4 +64,4 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test", "SafeTestsets", "DifferentialEquations", "MLJLinearModels", "LIBSVM", "Statistics", "SparseArrays", "CellularAutomata"]
+test = ["Aqua", "Test", "SafeTestsets", "DifferentialEquations", "MLJLinearModels", "LIBSVM", "Statistics", "SparseArrays", "CellularAutomata", "JET"]

ext/RCLIBSVMExt.jl

@@ -52,9 +52,9 @@ function (svmro::SVMReadout)(inp::AbstractArray, ps, st::NamedTuple)
 
     if models isa AbstractVector
         out_data = Array{float(eltype(reshaped_inp))}(undef, svmro.out_dims, num_imp)
-        @inbounds for i in 1:svmro.out_dims
-            single_out = LIBSVM.predict(models[i], reshaped_inp)
-            out_data[i, :] = single_out
+        for (idx, model) in enumerate(models)
+            single_out = LIBSVM.predict(models[idx], reshaped_inp)
+            out_data[idx, :] = single_out
         end
     else
         single_out = LIBSVM.predict(models, reshaped_inp)

src/ReservoirComputing.jl

@@ -4,7 +4,7 @@ using Adapt: adapt
 using ArrayInterface: ArrayInterface
 using Compat: @compat
 using ConcreteStructs: @concrete
-using LinearAlgebra: eigvals, mul!, I, qr, Diagonal
+using LinearAlgebra: eigvals, mul!, I, qr, Diagonal, diag
 using LuxCore: AbstractLuxLayer, AbstractLuxContainerLayer, AbstractLuxWrapperLayer,
     setup, apply, replicate
 import LuxCore: initialparameters, initialstates, statelength, outputsize

src/inits/esn_inits.jl

@@ -52,33 +52,27 @@ before this layer (logically inserting a [`Collect()`](@ref) right before it).
     activation
     in_dims <: IntegerType
     out_dims <: IntegerType
+    init_weight
+    init_bias
     use_bias <: StaticBool
     include_collect <: StaticBool
 end
 
-function Base.show(io::IO, ro::Readout)
-    print(io, "Readout($(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)))
-    ic === true && print(io, ", include_collect=true")
-    return print(io, ")")
-end
-
 function Readout(mapping::Pair{<:IntegerType,<:IntegerType}, activation=identity; kwargs...)
     return Readout(first(mapping), last(mapping), activation; kwargs...)
 end
 
 function Readout(in_dims::IntegerType, out_dims::IntegerType, activation=identity;
-    include_collect::BoolType=True(), use_bias::BoolType=False())
-    return Readout(activation, in_dims, out_dims, static(use_bias), static(include_collect))
+    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))
 end
 
 function initialparameters(rng::AbstractRNG, ro::Readout)
-    weight = rand(rng, Float32, ro.out_dims, ro.in_dims)
+    weight = ro.init_weight(rng, ro.out_dims, ro.in_dims)
 
     if has_bias(ro)
-        return (; weight, bias=rand(rng, Float32, ro.out_dims))
+        return (; weight, bias=ro.init_bias(rng, Float32, ro.out_dims))
     else
         return (; weight)
     end
@@ -98,6 +92,15 @@ 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)")
+    (ro.activation == identity) || print(io, ", $(ro.activation)")
+    has_bias(ro) || print(io, ", use_bias=false")
+    ic = known(getproperty(ro, Val(:include_collect)))
+    ic === true && print(io, ", include_collect=true")
+    return print(io, ")")
+end
+
 @doc raw"""
     Collect()
 
@@ -188,7 +191,7 @@ that time step.
 - `st`: Updated model states.
 
 """
-function collectstates(rc::AbstractLuxLayer, data::AbstractArray, ps, st::NamedTuple)
+function collectstates(rc::AbstractLuxLayer, data::AbstractMatrix, ps, st::NamedTuple)
     newst = st
     collected = Any[]
     for inp in eachcol(data)
@@ -209,3 +212,7 @@ function collectstates(rc::AbstractLuxLayer, data::AbstractArray, ps, st::NamedT
     states = eltype(data).(reduce(hcat, collected))
     return states, newst
 end
+
+function collectstates(rc::AbstractLuxLayer, data::AbstractVector, ps, st::NamedTuple)
+    return collectstates(rc, reshape(data, :, 1), ps, st)
+end

src/layers/basic.jl

@@ -75,7 +75,7 @@ Created by `initialstates(rng, esn)`:
     use_bias <: StaticBool
 end
 
-function ESNCell((in_dims, out_dims)::Pair{<:Int,<:Int}, activation=tanh;
+function ESNCell((in_dims, out_dims)::Pair{<:IntegerType,<:IntegerType}, activation=tanh;
     use_bias::BoolType=False(), init_bias=zeros32, init_reservoir=rand_sparse,
     init_input=weighted_init, init_state=randn32, leak_coefficient=1.0)
     return ESNCell(activation, in_dims, out_dims, init_bias, init_reservoir,

src/layers/esn_cell.jl

@@ -124,6 +124,10 @@ function named_tuple_layers(layers::Vararg{AbstractLuxLayer,N}) where {N}
     return NamedTuple{ntuple(i -> Symbol(:layer_, i), N)}(layers)
 end
 
+function index_namedtuple(nt::NamedTuple{fields}, idxs::AbstractArray) where {fields}
+    return NamedTuple{fields[idxs]}(values(nt)[idxs])
+end
+
 # from Lux extended_ops
 const KnownSymbolType{v} = Union{Val{v},StaticSymbol{v}}
 
@@ -132,6 +136,10 @@ function has_bias(l::AbstractLuxLayer)
     return ifelse(res === nothing, false, res)
 end
 
-function getproperty(x, ::KnownSymbolType{v}) where {v}
-    return v ∈ Base.propertynames(x) ? Base.getproperty(x, v) : nothing
+@generated function getproperty(x::X, ::KnownSymbolType{v}) where {X,v}
+    if hasfield(X, v)
+        return :(getfield(x, v))
+    else
+        return :(nothing)
+    end
 end

src/layers/lux_layers.jl

@@ -1,7 +1,13 @@
 
-function _apply_tomatrix(res_states, x_old::AbstractMatrix)
-    results = res_states.(eachcol(x_old))
-    return hcat(results...)
+@inline function _apply_tomatrix(states_mod::F, states::AbstractMatrix) where {F<:Function}
+    cols = axes(states, 2)
+    states_1 = states_mod(states[:, first(cols)])
+    new_states = similar(states_1, length(states_1), length(cols))
+    new_states[:, 1] .= states_1
+    for (k, j) in enumerate(cols)
+        new_states[:, k] .= states_mod(@view states[:, j])
+    end
+    return new_states
 end