fix: lower bound for JET

Author: MartinuzziFrancesco

Project.toml

@@ -34,7 +34,7 @@ CellularAutomata = "0.0.6"
 Compat = "4.16.0"
 ConcreteStructs = "0.2.3"
 DifferentialEquations = "7.16.1"
-JET = "0.10.9"
+JET = "0.9, 0.10.10"
 LIBSVM = "0.8"
 LinearAlgebra = "1.10"
 LuxCore = "1.3.0"

src/layers/lux_layers.jl

@@ -82,7 +82,9 @@ end
 (c::ReservoirChain)(x, ps, st::NamedTuple) = applychain(c.layers, x, ps, st)
 
 @generated function applychain(
-        layers::NamedTuple{fields}, x, ps, st::NamedTuple{fields}) where {fields}
+    layers::NamedTuple{fields}, x, ps, st::NamedTuple{fields}
+) where {fields}
+    @assert isa(fields, NTuple{<:Any, Symbol})
     N = length(fields)
     x_symbols = vcat([:x], [gensym() for _ in 1:N])
     st_symbols = [gensym() for _ in 1:N]

src/models/esn_generics.jl

@@ -12,23 +12,42 @@ _wrap_layers(xs::Tuple) = map(_wrap_layer, xs)
     return inp, tuple(new_st_parts...)
 end
 
-function _asvec(comp, n_layers::Integer)
-    if comp === ()
-        return ntuple(_ -> nothing, n_layers)
-    elseif comp isa Tuple || comp isa AbstractVector
-        len = length(comp)
-        if len == n_layers
-            return Tuple(comp)
-        elseif len == 1
-            return ntuple(_ -> comp[1], n_layers)
-        else
-            error("Expected length $n_layers or 1, got $len")
-        end
+@inline function _fillvec(x, n::Integer)
+    v = Vector{typeof(x)}(undef, n)
+    @inbounds @simd for i in 1:n
+        v[i] = x
+    end
+    return v
+end
+
+@inline _asvec(::Tuple{}, n::Integer) = _fillvec(nothing, n)
+
+@inline function _asvec(comp::Tuple, n::Integer)
+    len = length(comp)
+    if len == n
+        return collect(comp)
+    elseif len == 1
+        return _fillvec(comp[1], n)
+    else
+        error("Expected length $n or 1, got $len")
+    end
+end
+
+@inline function _asvec(comp::AbstractVector, n::Integer)
+    len = length(comp)
+    if len == n
+        return collect(comp)
+    elseif len == 1
+        return _fillvec(comp[1], n)
     else
-        return ntuple(_ -> comp, n_layers)
+        error("Expected length $n or 1, got $len")
     end
 end
 
+@inline _asvec(::Nothing, n::Integer) = _fillvec(nothing, n)
+
+@inline _asvec(comp, n::Integer) = _fillvec(comp, n)
+
 @inline _asvec(x) = (ndims(x) == 2 ? vec(x) : x)
 
 function _coerce_layer_mods(x)