From 9f1575f7f7a804a44c4e4c6f52242f90bc9dc875 Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Tue, 15 Sep 2026 12:23:57 +0530 Subject: [PATCH 1/2] feat: closed-loop ESN Jacobian API (#169) Add jacobian/jacobian!/jacobians for the autonomous reservoir map of a trained ESN (analytical default; ForwardDiff weakdep fallback), with Models tests and API docs. --- Project.toml | 6 +- docs/pages.jl | 1 + docs/src/api/jacobian.md | 7 + ext/RCForwardDiffExt.jl | 17 ++ src/ReservoirComputing.jl | 5 +- src/jacobian.jl | 414 ++++++++++++++++++++++++++++++++++ test/Models/jacobian_tests.jl | 185 +++++++++++++++ test/qa/qa.jl | 10 +- 8 files changed, 638 insertions(+), 7 deletions(-) create mode 100644 docs/src/api/jacobian.md create mode 100644 ext/RCForwardDiffExt.jl create mode 100644 src/jacobian.jl create mode 100644 test/Models/jacobian_tests.jl diff --git a/Project.toml b/Project.toml index 5b1e8c53e..941d5137d 100644 --- a/Project.toml +++ b/Project.toml @@ -19,6 +19,7 @@ WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d" [weakdeps] CellularAutomata = "878138dc-5b27-11ea-1a71-cb95d38d6b29" DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LIBSVM = "b1bec4e5-fd48-53fe-b0cb-9723c09d164b" MLJLinearModels = "6ee0df7b-362f-4a72-a706-9e79364fb692" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" @@ -26,6 +27,7 @@ StateSpaceSets = "40b095a5-5852-4c12-98c7-d43bf788e795" [extensions] RCCellularAutomataExt = "CellularAutomata" +RCForwardDiffExt = "ForwardDiff" RCLIBSVMExt = "LIBSVM" RCMLJLinearModelsExt = "MLJLinearModels" RCODEReservoirExt = "DataInterpolations" @@ -39,6 +41,7 @@ CellularAutomata = "0.0.6, 0.1" ConcreteStructs = "0.2.3" DataInterpolations = "9.0, 10.1" DifferentialEquations = "8" +ForwardDiff = "0.10, 1" LIBSVM = "0.8" LinearAlgebra = "1.10" LinearSolve = "5.1" @@ -68,6 +71,7 @@ julia = "1.10" CellularAutomata = "878138dc-5b27-11ea-1a71-cb95d38d6b29" DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LIBSVM = "b1bec4e5-fd48-53fe-b0cb-9723c09d164b" MLJLinearModels = "6ee0df7b-362f-4a72-a706-9e79364fb692" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" @@ -83,4 +87,4 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test", "SafeTestsets", "SciMLTesting", "CellularAutomata", "DataInterpolations", "DifferentialEquations", "MLJLinearModels", "LIBSVM", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqLowOrderRK", "Serialization", "SparseArrays", "StateSpaceSets", "StaticArrays", "Statistics"] +test = ["Test", "SafeTestsets", "SciMLTesting", "CellularAutomata", "DataInterpolations", "DifferentialEquations", "ForwardDiff", "MLJLinearModels", "LIBSVM", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqLowOrderRK", "Serialization", "SparseArrays", "StateSpaceSets", "StaticArrays", "Statistics"] diff --git a/docs/pages.jl b/docs/pages.jl index 935e32ff3..5fe9a99b9 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -27,6 +27,7 @@ pages = [ "Utilities" => "api/utils.md", "Train" => "api/train.md", "Predict" => "api/predict.md", + "Jacobian" => "api/jacobian.md", "States" => "api/states.md", "Initializers" => "api/inits.md", "Developer Interfaces" => "api/developer.md", diff --git a/docs/src/api/jacobian.md b/docs/src/api/jacobian.md new file mode 100644 index 000000000..1139846aa --- /dev/null +++ b/docs/src/api/jacobian.md @@ -0,0 +1,7 @@ +# Jacobian + +```@docs + jacobian + jacobian! + jacobians +``` diff --git a/ext/RCForwardDiffExt.jl b/ext/RCForwardDiffExt.jl new file mode 100644 index 000000000..5e9185b60 --- /dev/null +++ b/ext/RCForwardDiffExt.jl @@ -0,0 +1,17 @@ +module RCForwardDiffExt + +using ReservoirComputing: ReservoirComputing +using ForwardDiff: ForwardDiff + +function forwarddiff_closedloop_jacobian!( + J::AbstractMatrix, esn::ReservoirComputing.ESN, x::AbstractVector, ps + ) + ReservoirComputing.__require_esn_closedloop_io(esn) + closed_loop = let esn = esn, ps = ps + state -> first(ReservoirComputing.__closed_loop_step(esn, state, ps)) + end + J .= ForwardDiff.jacobian(closed_loop, x) + return J +end + +end # module diff --git a/src/ReservoirComputing.jl b/src/ReservoirComputing.jl index 44b93d7ad..3618a57df 100644 --- a/src/ReservoirComputing.jl +++ b/src/ReservoirComputing.jl @@ -68,6 +68,7 @@ include("models/rmnesn.jl") include("models/rmnresesn.jl") include("models/continuous_esn.jl") include("models/lsm.jl") +include("jacobian.jl") #conceptors include("conceptors.jl") #extensions @@ -109,8 +110,8 @@ export band_init, block_diagonal, chaotic_init, cycle_jumps, delay_line, delayli export add_jumps!, backward_connection!, delay_line!, permute_matrix!, reverse_simple_cycle!, scale_radius!, self_loop!, simple_cycle! export @topology -export polynomial_monomials, chebyshev_monomials, predict, QRSolver, QRFactorization, - resetcarry!, return_init_as, train, train! +export polynomial_monomials, chebyshev_monomials, predict, jacobian, jacobian!, jacobians, + QRSolver, QRFactorization, resetcarry!, return_init_as, train, train! export AdditiveEIESN, DeepESN, DeepReservoir, DelayESN, EIESN, ES2N, ESN, EuSN, HybridESN, InputDelayESN, LIFESN, ResESN, StateDelayESN, SVESM export NGRC export RMNESN, RMNResESN diff --git a/src/jacobian.jl b/src/jacobian.jl new file mode 100644 index 000000000..59c1ff1df --- /dev/null +++ b/src/jacobian.jl @@ -0,0 +1,414 @@ +@doc raw""" + jacobian(esn::ESN, state, ps, st; backend=:analytical) + jacobian!(J, esn::ESN, state, ps, st; backend=:analytical) + +Jacobian of the closed-loop reservoir map at reservoir state `state` +[Pathak2017](@cite). + +Requires `in_dims == out_dims`. With readout feedback the reservoir obeys + +```math +\begin{aligned} +\mathbf{z} &= \mathrm{Mods}(\mathbf{x}), \\ +\mathbf{u} &= \rho(\mathbf{W}_{\mathrm{out}}\mathbf{z}+\mathbf{b}_{\mathrm{out}}), \\ +\mathbf{a} &= \mathbf{W}_{\mathrm{in}}\mathbf{u}+\mathbf{W}_r\mathbf{x}+\mathbf{b}, \\ +F(\mathbf{x}) &= (1-\alpha)\odot\mathbf{x}+\alpha\odot\varphi(\mathbf{a}), +\end{aligned} +``` + +and this returns ``J = DF/D\mathbf{x}``. The carry is not advanced. + +## Arguments + + - `esn`: an [`ESN`](@ref). + - `state`: reservoir carry, length `res_dims` (vector or single-column matrix). + - `ps`: model parameters. + - `st`: model states. + - `J`: preallocated `res_dims × res_dims` buffer for `jacobian!`. + +## Keyword arguments + + - `backend`: `:analytical` (default) or `:forwarddiff`. + +## Returns + + - `(J, st)` with `J` of size `(res_dims, res_dims)`. `st` is unchanged. +""" +function jacobian( + esn::ESN, state::AbstractVecOrMat, ps, st; + backend::Symbol = :analytical + ) + x = __jacobian_state_vector(state) + J = Matrix{eltype(x)}(undef, length(x), length(x)) + return jacobian!(J, esn, x, ps, st; backend) +end + +function jacobian!( + J::AbstractMatrix, esn::ESN, state::AbstractVecOrMat, ps, st; + backend::Symbol = :analytical + ) + x = __jacobian_state_vector(state) + __check_closedloop_jacobian_dims(esn, x, J) + if backend === :analytical + __analytical_closedloop_jacobian!(J, esn, x, ps) + elseif backend === :forwarddiff + __forwarddiff_closedloop_jacobian!(J, esn, x, ps) + else + throw(ArgumentError("backend must be :analytical or :forwarddiff, got $(repr(backend))")) + end + return J, st +end + +@doc raw""" + jacobians(esn::ESN, steps, ps, st; initialdata, backend=:analytical) + +Jacobians of the closed-loop reservoir map along an autoregressive trajectory. + +Uses the same feedback as [`predict`](@ref): the first step is driven by +`initialdata`, then each output is fed back as the next input. After step +``t``, `Js[:, :, t]` stores ``DF(\mathbf{x}_t)`` at the current carry. + +## Arguments + + - `esn`: an [`ESN`](@ref) with `in_dims == out_dims`. + - `steps`: number of autoregressive steps. + - `ps`: model parameters. + - `st`: model states. + +## Keyword arguments + + - `initialdata`: column vector used as the first input. + - `backend`: `:analytical` (default) or `:forwarddiff`. + +## Returns + + - `Js`: array of size `(res_dims, res_dims, steps)`. + - `outputs`: generated outputs of shape `(out_dims, steps)`. + - `st`: model state after `steps` updates. +""" +function jacobians( + esn::ESN, steps::Integer, ps, st; + initialdata::AbstractVector, backend::Symbol = :analytical + ) + steps ≥ 1 || throw(ArgumentError("steps must be ≥ 1, got $steps")) + input_length = length(initialdata) + __require_esn_closedloop_io(esn) + + current_output, st = apply(esn, initialdata, ps, st) + __require_closed_loop_dimension(current_output, input_length, 1) + + x = __carry_state_vector(st) + n = length(x) + Js = Array{eltype(x)}(undef, n, n, steps) + outputs = similar(current_output, length(current_output), steps) + + jacobian!(view(Js, :, :, 1), esn, x, ps, st; backend) + outputs[:, 1] .= current_output + + for step in 2:steps + current_output, st = apply(esn, current_output, ps, st) + __require_closed_loop_dimension(current_output, input_length, step) + x = __carry_state_vector(st) + jacobian!(view(Js, :, :, step), esn, x, ps, st; backend) + outputs[:, step] .= current_output + end + return Js, outputs, st +end + +function __jacobian_state_vector(state::AbstractVector) + return state +end + +function __jacobian_state_vector(state::AbstractMatrix) + size(state, 2) == 1 || throw( + ArgumentError( + "jacobian expects a vector or single-column matrix, got size $(size(state))" + ) + ) + return vec(state) +end + +function __carry_state_vector(st::NamedTuple) + carry = get(st.reservoir, :carry, nothing) + carry === nothing && throw( + ArgumentError("reservoir carry is unset; run a model step or pass `state` explicitly") + ) + return __jacobian_state_vector(first(carry)) +end + +function __require_esn_closedloop_io(esn::ESN) + in_dims = Int(esn.reservoir.cell.in_dims) + out_dims = Int(esn.readout.out_dims) + in_dims == out_dims || throw( + DimensionMismatch( + "closed-loop jacobian requires in_dims == out_dims (got $in_dims and $out_dims)" + ) + ) + return nothing +end + +function __check_closedloop_jacobian_dims(esn::ESN, x::AbstractVector, J::AbstractMatrix) + __require_esn_closedloop_io(esn) + n = Int(esn.reservoir.cell.out_dims) + length(x) == n || throw(DimensionMismatch("reservoir state length must be $n, got $(length(x))")) + size(J) == (n, n) || throw(DimensionMismatch("Jacobian buffer must be ($n, $n), got $(size(J))")) + return nothing +end + +function __closed_loop_quantities(esn::ESN, x::AbstractVector, ps) + cell = esn.reservoir.cell + z = __apply_modifiers_pure(esn.state_modifiers, x, ps.state_modifiers) + u = __readout_pure(esn.readout, z, ps.readout) + + input_matrix = ps.reservoir.input_matrix + reservoir_matrix = ps.reservoir.reservoir_matrix + bias = safe_getproperty(ps.reservoir, Val(:bias)) + preactivation = dense_bias(input_matrix, u, nothing) .+ + dense_bias(reservoir_matrix, x, bias) + + T = eltype(x) + leak = __format_leak(T, cell.leak_coefficient) + x_new = __one_minus_leak(T, leak) .* x .+ leak .* cell.activation.(preactivation) + return (; x_new, u, preactivation, z, leak, input_matrix, reservoir_matrix) +end + +function __closed_loop_step(esn::ESN, x::AbstractVector, ps) + q = __closed_loop_quantities(esn, x, ps) + return q.x_new, q.u +end + +__apply_modifiers_pure(::Tuple{}, x, ::Tuple{}) = x + +function __apply_modifiers_pure(modifiers::Tuple, x, ps_mods::Tuple) + features = x + for (modifier, ps_mod) in zip(modifiers, ps_mods) + features, _ = __apply_state_modifier(modifier, features, x, ps_mod, NamedTuple()) + end + return features +end + +function __readout_pure(readout::LinearReadout, features, ps_readout) + out = ps_readout.weight * features + if has_bias(readout) + out = out .+ ps_readout.bias + end + return readout.activation.(out) +end + +function __analytical_closedloop_jacobian!(J::AbstractMatrix, esn::ESN, x::AbstractVector, ps) + __ensure_supported_modifiers(esn.state_modifiers) + q = __closed_loop_quantities(esn, x, ps) + M = __state_modifiers_jacobian(esn.state_modifiers, x) + Du_Dx = __readout_jacobian(esn.readout, q.z, ps.readout, M) + + J .= q.input_matrix * Du_Dx + J .+= q.reservoir_matrix # densifies if `reservoir_matrix` is sparse + + dφ = __activation_derivative(esn.reservoir.cell.activation, q.preactivation) + __scale_rows!(J, q.leak, dφ) + __add_leak_identity!(J, q.leak) + return J +end + +function __scale_rows!(J::AbstractMatrix, leak::Number, dφ::AbstractVector) + @inbounds for i in axes(J, 1) + scale = leak * dφ[i] + for j in axes(J, 2) + J[i, j] *= scale + end + end + return J +end + +function __scale_rows!(J::AbstractMatrix, leak::AbstractArray, dφ::AbstractVector) + leak_vec = vec(leak) + length(leak_vec) == length(dφ) || throw( + DimensionMismatch( + "leak_coefficient length $(length(leak_vec)) must match reservoir size $(length(dφ))" + ) + ) + @inbounds for i in axes(J, 1) + scale = leak_vec[i] * dφ[i] + for j in axes(J, 2) + J[i, j] *= scale + end + end + return J +end + +function __add_leak_identity!(J::AbstractMatrix, leak::Number) + one_minus = one(eltype(J)) - leak + @inbounds for i in axes(J, 1) + J[i, i] += one_minus + end + return J +end + +function __add_leak_identity!(J::AbstractMatrix, leak::AbstractArray) + leak_vec = vec(leak) + @inbounds for i in axes(J, 1) + J[i, i] += one(eltype(J)) - leak_vec[i] + end + return J +end + +function __readout_jacobian(readout::LinearReadout, z, ps_readout, M::AbstractMatrix) + weight_M = ps_readout.weight * M + readout.activation === identity && return weight_M + + pre = ps_readout.weight * z + if has_bias(readout) + pre = pre .+ ps_readout.bias + end + return __activation_derivative(readout.activation, pre) .* weight_M +end + +__activation_derivative(::typeof(identity), a::AbstractVector) = ones(eltype(a), length(a)) + +function __activation_derivative(::typeof(tanh), a::AbstractVector) + y = tanh.(a) + return one(eltype(a)) .- y .* y +end + +function __activation_derivative(::typeof(tanh_fast), a::AbstractVector) + y = tanh_fast.(a) + return one(eltype(a)) .- y .* y +end + +function __activation_derivative(activation, ::AbstractVector) + throw( + ArgumentError( + "no analytical derivative for activation $(activation); " * + "use `backend=:forwarddiff` after `using ForwardDiff`" + ) + ) +end + +__unwrap_modifier(wf::WrappedFunction) = wf.func +__unwrap_modifier(modifier) = modifier + +function __ensure_supported_modifiers(modifiers::Tuple) + for modifier in modifiers + unwrapped = __unwrap_modifier(modifier) + if unwrapped isa Extend + throw( + ArgumentError( + "closed-loop jacobian does not support `Extend` " * + "(autonomous state is not the reservoir carry alone)" + ) + ) + end + if !hasmethod(__state_modifier_jacobian, Tuple{typeof(unwrapped), AbstractVector}) + throw( + ArgumentError( + "no analytical Jacobian for state modifier $(unwrapped); " * + "use `backend=:forwarddiff` after `using ForwardDiff`" + ) + ) + end + end + return nothing +end + +function __state_modifiers_jacobian(::Tuple{}, x::AbstractVector) + n = length(x) + return Matrix{eltype(x)}(I, n, n) +end + +function __state_modifiers_jacobian(modifiers::Tuple, x::AbstractVector) + features = x + unwrapped = __unwrap_modifier(first(modifiers)) + M = __state_modifier_jacobian(unwrapped, features) + features = unwrapped(features) + for modifier in Base.tail(modifiers) + unwrapped = __unwrap_modifier(modifier) + M = __state_modifier_jacobian(unwrapped, features) * M + features = unwrapped(features) + end + return M +end + +function __state_modifier_jacobian(::typeof(NLAT1), x::AbstractVector) + T = eltype(x) + n = length(x) + M = Matrix{T}(I, n, n) + @inbounds for i in eachindex(x) + if isodd(i) + M[i, i] = 2 * x[i] + end + end + return M +end + +function __state_modifier_jacobian(::typeof(NLAT2), x::AbstractVector) + T = eltype(x) + n = length(x) + M = Matrix{T}(I, n, n) + first_i = firstindex(x) + @inbounds for i in eachindex(x) + if i > first_i && isodd(i) + M[i, i] = zero(T) + M[i, i - 1] = x[i - 2] + M[i, i - 2] = x[i - 1] + end + end + return M +end + +function __state_modifier_jacobian(::typeof(NLAT3), x::AbstractVector) + T = eltype(x) + n = length(x) + M = Matrix{T}(I, n, n) + first_i = firstindex(x) + last_i = lastindex(x) + @inbounds for i in eachindex(x) + if first_i < i < last_i && isodd(i) + M[i, i] = zero(T) + M[i, i - 1] = x[i + 1] + M[i, i + 1] = x[i - 1] + end + end + return M +end + +function __state_modifier_jacobian(::Pad, x::AbstractVector) + T = eltype(x) + n = length(x) + M = zeros(T, n + 1, n) + @inbounds for i in 1:n + M[i, i] = one(T) + end + return M +end + +function __state_modifier_jacobian(partial_square::PartialSquare, x::AbstractVector) + T = eltype(x) + n = length(x) + M = Matrix{T}(I, n, n) + threshold = floor(Int, partial_square.eta * n) + @inbounds for i in 1:threshold + M[i, i] = 2 * x[i] + end + return M +end + +function __state_modifier_jacobian(::typeof(ExtendedSquare), x::AbstractVector) + T = eltype(x) + n = length(x) + M = zeros(T, 2n, n) + @inbounds for i in 1:n + M[i, i] = one(T) + M[n + i, i] = 2 * x[i] + end + return M +end + +function __forwarddiff_closedloop_jacobian!( + J::AbstractMatrix, esn::ESN, x::AbstractVector, ps + ) + ext = Base.get_extension(@__MODULE__, :RCForwardDiffExt) + ext === nothing && error( + "backend=:forwarddiff requires ForwardDiff (`using ForwardDiff`)" + ) + return ext.forwarddiff_closedloop_jacobian!(J, esn, x, ps) +end diff --git a/test/Models/jacobian_tests.jl b/test/Models/jacobian_tests.jl new file mode 100644 index 000000000..779febce5 --- /dev/null +++ b/test/Models/jacobian_tests.jl @@ -0,0 +1,185 @@ +# Closed-loop ESN Jacobian. +begin + using Test + using Random + using LinearAlgebra + using ReservoirComputing + using LuxCore: setup + using ForwardDiff + using NNlib: tanh_fast + + const _ATOL = 5.0f-3 + + function _finite_diff_jacobian(esn, state, ps; ε = 1.0f-3) + n = length(state) + J = zeros(eltype(state), n, n) + for j in 1:n + state_plus = copy(state) + state_minus = copy(state) + state_plus[j] += ε + state_minus[j] -= ε + forward_plus = first(ReservoirComputing.__closed_loop_step(esn, state_plus, ps)) + forward_minus = first(ReservoirComputing.__closed_loop_step(esn, state_minus, ps)) + J[:, j] .= (forward_plus .- forward_minus) ./ (2 * ε) + end + return J + end + + function _with_readout_weights(esn, ps, rng) + out_dims = Int(esn.readout.out_dims) + feat_dims = Int(esn.readout.in_dims) + weight = randn(rng, Float32, out_dims, feat_dims) .* 0.05f0 + if haskey(ps.readout, :bias) + return merge(ps, (readout = (weight = weight, bias = zeros(Float32, out_dims)),)) + end + return merge(ps, (readout = (weight = weight,),)) + end + + @testset "analytical matches finite differences" begin + rng = MersenneTwister(42) + esn = ESN( + 3, 6, 3; + init_reservoir = scaled_rand, + use_bias = true, + leak_coefficient = 0.7f0, + ) + ps, st = setup(rng, esn) + ps = _with_readout_weights(esn, ps, rng) + state = randn(rng, Float32, 6) + + J_analytical, st_out = jacobian(esn, state, ps, st) + @test st_out === st + @test size(J_analytical) == (6, 6) + @test eltype(J_analytical) === Float32 + @test J_analytical ≈ _finite_diff_jacobian(esn, state, ps) atol = _ATOL + end + + @testset "ForwardDiff backend" begin + rng = MersenneTwister(7) + esn = ESN(2, 5, 2; init_reservoir = scaled_rand) + ps, st = setup(rng, esn) + ps = _with_readout_weights(esn, ps, rng) + state = randn(rng, Float32, 5) + J_analytical, _ = jacobian(esn, state, ps, st) + J_ad, _ = jacobian(esn, state, ps, st; backend = :forwarddiff) + @test J_analytical ≈ J_ad atol = 1.0f-5 + end + + @testset "vector leak" begin + rng = MersenneTwister(11) + leak = Float32[0.5, 0.7, 0.9, 0.6, 0.8] + esn = ESN(2, 5, 2; init_reservoir = scaled_rand, leak_coefficient = leak) + ps, st = setup(rng, esn) + ps = _with_readout_weights(esn, ps, rng) + state = randn(rng, Float32, 5) + J_analytical, _ = jacobian(esn, state, ps, st) + @test J_analytical ≈ _finite_diff_jacobian(esn, state, ps) atol = _ATOL + end + + @testset "state modifiers" begin + rng = MersenneTwister(13) + for modifier in (NLAT1, NLAT2, NLAT3, PartialSquare(0.5), ExtendedSquare, Pad(1.0f0)) + res_dims = modifier isa Pad ? 5 : 6 + readout_in_dims = if modifier === ExtendedSquare + 12 + elseif modifier isa Pad + 6 + else + res_dims + end + esn = ESN( + 2, res_dims, 2; + init_reservoir = scaled_rand, + state_modifiers = modifier, + readout_in_dims = readout_in_dims, + ) + ps, st = setup(rng, esn) + ps = _with_readout_weights(esn, ps, rng) + state = randn(rng, Float32, res_dims) + J_analytical, _ = jacobian(esn, state, ps, st) + @test J_analytical ≈ _finite_diff_jacobian(esn, state, ps) atol = _ATOL + end + end + + @testset "NLAT2 sparsity pattern" begin + state = Float32[1, 2, 3, 4, 5] + M = ReservoirComputing.__state_modifier_jacobian(NLAT2, state) + @test M[3, 3] == 0 + @test M[3, 2] == state[1] + @test M[3, 1] == state[2] + @test M[5, 5] == 0 + @test M[5, 4] == state[3] + @test M[5, 3] == state[4] + end + + @testset "tanh_fast" begin + rng = MersenneTwister(19) + esn = ESN(2, 5, 2, tanh_fast; init_reservoir = scaled_rand) + ps, st = setup(rng, esn) + ps = _with_readout_weights(esn, ps, rng) + state = randn(rng, Float32, 5) + J_analytical, _ = jacobian(esn, state, ps, st) + J_ad, _ = jacobian(esn, state, ps, st; backend = :forwarddiff) + @test J_analytical ≈ J_ad atol = 1.0f-5 + end + + @testset "jacobians matches predict" begin + rng = MersenneTwister(23) + esn = ESN(3, 6, 3; init_reservoir = scaled_rand) + ps, st = setup(rng, esn) + ps = _with_readout_weights(esn, ps, rng) + initialdata = Float32[0.1, -0.2, 0.3] + + Js, outputs, st_final = jacobians(esn, 4, ps, st; initialdata) + predicted, st_pred = predict(esn, 4, ps, st; initialdata) + @test outputs ≈ predicted + @test size(Js) == (6, 6, 4) + @test st_final.reservoir.carry == st_pred.reservoir.carry + + _, st_step = apply(esn, initialdata, ps, st) + for t in 1:4 + state = ReservoirComputing.__carry_state_vector(st_step) + Jt, _ = jacobian(esn, state, ps, st_step) + @test Js[:, :, t] ≈ Jt atol = 1.0f-6 + if t < 4 + _, st_step = apply(esn, outputs[:, t], ps, st_step) + end + end + end + + @testset "errors" begin + rng = MersenneTwister(29) + esn_bad = ESN(3, 5, 2; init_reservoir = scaled_rand) + ps, st = setup(rng, esn_bad) + state = randn(rng, Float32, 5) + @test_throws DimensionMismatch jacobian(esn_bad, state, ps, st) + + esn_ext = ESN( + 3, 5, 3; + init_reservoir = scaled_rand, + state_modifiers = Extend(Collect()), + readout_in_dims = 8, + ) + ps_ext, st_ext = setup(rng, esn_ext) + @test_throws ArgumentError jacobian(esn_ext, state, ps_ext, st_ext) + + esn = ESN(3, 5, 3; init_reservoir = scaled_rand) + ps, st = setup(rng, esn) + @test_throws ArgumentError jacobian(esn, state, ps, st; backend = :something) + @test_throws DimensionMismatch jacobian!(zeros(Float32, 4, 4), esn, state, ps, st) + @test_throws ArgumentError jacobians(esn, 0, ps, st; initialdata = Float32[1, 2, 3]) + end + + @testset "jacobian!" begin + rng = MersenneTwister(31) + esn = ESN(2, 4, 2; init_reservoir = scaled_rand) + ps, st = setup(rng, esn) + ps = _with_readout_weights(esn, ps, rng) + state = randn(rng, Float32, 4) + J = zeros(Float32, 4, 4) + J_out, _ = jacobian!(J, esn, state, ps, st) + @test J_out === J + J_ref, _ = jacobian(esn, state, ps, st) + @test J ≈ J_ref + end +end diff --git a/test/qa/qa.jl b/test/qa/qa.jl index 1e38a5192..e83fc1dd1 100644 --- a/test/qa/qa.jl +++ b/test/qa/qa.jl @@ -3,10 +3,10 @@ using JET # ExplicitImports only checks an extension module once it exists, and an extension only # exists once its trigger package is loaded. Loading the weakdeps here is what puts -# RCCellularAutomataExt, RCODEReservoirExt, RCLIBSVMExt, RCMLJLinearModelsExt, -# RCSparseArraysExt and RCStateSpaceSetsExt in scope for the QA checks. -using CellularAutomata, DataInterpolations, LIBSVM, MLJLinearModels, SparseArrays, - StateSpaceSets +# RCCellularAutomataExt, RCForwardDiffExt, RCODEReservoirExt, RCLIBSVMExt, +# RCMLJLinearModelsExt, RCSparseArraysExt and RCStateSpaceSetsExt in scope for the QA checks. +using CellularAutomata, DataInterpolations, ForwardDiff, LIBSVM, MLJLinearModels, + SparseArrays, StateSpaceSets # ReservoirComputing's own extension hook points. ExplicitImports' `allow_internal_imports` # / `allow_internal_accesses` defaults would cover these, but they key off @@ -24,10 +24,12 @@ rc_internal_hooks = ( :__check_protected_kwargs, :__collectstates, :__continuous_esn_rhs!, + :__closed_loop_step, :__feature_dim, :__fit_readout, :__init_encoder_st, :__predict, + :__require_esn_closedloop_io, :__reservoir_jac_prototype, :__resolve_readout_in_dims, :__supports_ar, From 752c3110c1af181cd76c99d3075c05335578027a Mon Sep 17 00:00:00 2001 From: Saswat Susmoy Date: Tue, 15 Sep 2026 12:28:57 +0530 Subject: [PATCH 2/2] refactor: share leak scaling path in closed-loop Jacobian Collapse row-scale/leak update into one finalize step, reuse LinearReadout for the pure closed-loop map, and drive jacobians through a single rollout loop shared with predict semantics. --- src/jacobian.jl | 189 ++++++++++------------------------ test/Models/jacobian_tests.jl | 115 +++++++-------------- 2 files changed, 94 insertions(+), 210 deletions(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index 59c1ff1df..6d829cdea 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -91,33 +91,28 @@ function jacobians( initialdata::AbstractVector, backend::Symbol = :analytical ) steps ≥ 1 || throw(ArgumentError("steps must be ≥ 1, got $steps")) - input_length = length(initialdata) __require_esn_closedloop_io(esn) - - current_output, st = apply(esn, initialdata, ps, st) - __require_closed_loop_dimension(current_output, input_length, 1) - - x = __carry_state_vector(st) - n = length(x) - Js = Array{eltype(x)}(undef, n, n, steps) - outputs = similar(current_output, length(current_output), steps) - - jacobian!(view(Js, :, :, 1), esn, x, ps, st; backend) - outputs[:, 1] .= current_output - - for step in 2:steps - current_output, st = apply(esn, current_output, ps, st) + input_length = length(initialdata) + current_input = initialdata + outputs = nothing + Js = nothing + for step in 1:steps + current_output, st = apply(esn, current_input, ps, st) __require_closed_loop_dimension(current_output, input_length, step) x = __carry_state_vector(st) + if step == 1 + n = length(x) + Js = Array{eltype(x)}(undef, n, n, steps) + outputs = similar(current_output, length(current_output), steps) + end jacobian!(view(Js, :, :, step), esn, x, ps, st; backend) outputs[:, step] .= current_output + current_input = current_output end return Js, outputs, st end -function __jacobian_state_vector(state::AbstractVector) - return state -end +__jacobian_state_vector(state::AbstractVector) = state function __jacobian_state_vector(state::AbstractMatrix) size(state, 2) == 1 || throw( @@ -158,24 +153,20 @@ end function __closed_loop_quantities(esn::ESN, x::AbstractVector, ps) cell = esn.reservoir.cell z = __apply_modifiers_pure(esn.state_modifiers, x, ps.state_modifiers) - u = __readout_pure(esn.readout, z, ps.readout) - + u = first(esn.readout(z, ps.readout, NamedTuple())) input_matrix = ps.reservoir.input_matrix reservoir_matrix = ps.reservoir.reservoir_matrix bias = safe_getproperty(ps.reservoir, Val(:bias)) preactivation = dense_bias(input_matrix, u, nothing) .+ dense_bias(reservoir_matrix, x, bias) - T = eltype(x) leak = __format_leak(T, cell.leak_coefficient) x_new = __one_minus_leak(T, leak) .* x .+ leak .* cell.activation.(preactivation) return (; x_new, u, preactivation, z, leak, input_matrix, reservoir_matrix) end -function __closed_loop_step(esn::ESN, x::AbstractVector, ps) - q = __closed_loop_quantities(esn, x, ps) - return q.x_new, q.u -end +__closed_loop_step(esn::ESN, x::AbstractVector, ps) = + ((q = __closed_loop_quantities(esn, x, ps)); (q.x_new, q.u)) __apply_modifiers_pure(::Tuple{}, x, ::Tuple{}) = x @@ -187,67 +178,35 @@ function __apply_modifiers_pure(modifiers::Tuple, x, ps_mods::Tuple) return features end -function __readout_pure(readout::LinearReadout, features, ps_readout) - out = ps_readout.weight * features - if has_bias(readout) - out = out .+ ps_readout.bias - end - return readout.activation.(out) -end - function __analytical_closedloop_jacobian!(J::AbstractMatrix, esn::ESN, x::AbstractVector, ps) __ensure_supported_modifiers(esn.state_modifiers) q = __closed_loop_quantities(esn, x, ps) M = __state_modifiers_jacobian(esn.state_modifiers, x) - Du_Dx = __readout_jacobian(esn.readout, q.z, ps.readout, M) - - J .= q.input_matrix * Du_Dx + J .= q.input_matrix * __readout_jacobian(esn.readout, q.z, ps.readout, M) J .+= q.reservoir_matrix # densifies if `reservoir_matrix` is sparse - - dφ = __activation_derivative(esn.reservoir.cell.activation, q.preactivation) - __scale_rows!(J, q.leak, dφ) - __add_leak_identity!(J, q.leak) - return J + return __finalize_leak_jacobian!( + J, q.leak, __activation_derivative(esn.reservoir.cell.activation, q.preactivation) + ) end -function __scale_rows!(J::AbstractMatrix, leak::Number, dφ::AbstractVector) +function __finalize_leak_jacobian!(J::AbstractMatrix, leak::Number, dφ::AbstractVector) + J .*= leak .* dφ @inbounds for i in axes(J, 1) - scale = leak * dφ[i] - for j in axes(J, 2) - J[i, j] *= scale - end + J[i, i] += one(eltype(J)) - leak end return J end -function __scale_rows!(J::AbstractMatrix, leak::AbstractArray, dφ::AbstractVector) - leak_vec = vec(leak) - length(leak_vec) == length(dφ) || throw( +function __finalize_leak_jacobian!(J::AbstractMatrix, leak::AbstractArray, dφ::AbstractVector) + α = vec(leak) + length(α) == length(dφ) || throw( DimensionMismatch( - "leak_coefficient length $(length(leak_vec)) must match reservoir size $(length(dφ))" + "leak_coefficient length $(length(α)) must match reservoir size $(length(dφ))" ) ) + J .*= α .* dφ @inbounds for i in axes(J, 1) - scale = leak_vec[i] * dφ[i] - for j in axes(J, 2) - J[i, j] *= scale - end - end - return J -end - -function __add_leak_identity!(J::AbstractMatrix, leak::Number) - one_minus = one(eltype(J)) - leak - @inbounds for i in axes(J, 1) - J[i, i] += one_minus - end - return J -end - -function __add_leak_identity!(J::AbstractMatrix, leak::AbstractArray) - leak_vec = vec(leak) - @inbounds for i in axes(J, 1) - J[i, i] += one(eltype(J)) - leak_vec[i] + J[i, i] += one(eltype(J)) - α[i] end return J end @@ -255,25 +214,15 @@ end function __readout_jacobian(readout::LinearReadout, z, ps_readout, M::AbstractMatrix) weight_M = ps_readout.weight * M readout.activation === identity && return weight_M - pre = ps_readout.weight * z - if has_bias(readout) - pre = pre .+ ps_readout.bias - end + has_bias(readout) && (pre = pre .+ ps_readout.bias) return __activation_derivative(readout.activation, pre) .* weight_M end __activation_derivative(::typeof(identity), a::AbstractVector) = ones(eltype(a), length(a)) - -function __activation_derivative(::typeof(tanh), a::AbstractVector) - y = tanh.(a) - return one(eltype(a)) .- y .* y -end - -function __activation_derivative(::typeof(tanh_fast), a::AbstractVector) - y = tanh_fast.(a) - return one(eltype(a)) .- y .* y -end +__activation_derivative(::typeof(tanh), a::AbstractVector) = (y = tanh.(a); one(eltype(a)) .- y .* y) +__activation_derivative(::typeof(tanh_fast), a::AbstractVector) = + (y = tanh_fast.(a); one(eltype(a)) .- y .* y) function __activation_derivative(activation, ::AbstractVector) throw( @@ -290,30 +239,25 @@ __unwrap_modifier(modifier) = modifier function __ensure_supported_modifiers(modifiers::Tuple) for modifier in modifiers unwrapped = __unwrap_modifier(modifier) - if unwrapped isa Extend - throw( - ArgumentError( - "closed-loop jacobian does not support `Extend` " * - "(autonomous state is not the reservoir carry alone)" - ) + unwrapped isa Extend && throw( + ArgumentError( + "closed-loop jacobian does not support `Extend` " * + "(autonomous state is not the reservoir carry alone)" ) - end - if !hasmethod(__state_modifier_jacobian, Tuple{typeof(unwrapped), AbstractVector}) + ) + hasmethod(__state_modifier_jacobian, Tuple{typeof(unwrapped), AbstractVector}) || throw( - ArgumentError( - "no analytical Jacobian for state modifier $(unwrapped); " * - "use `backend=:forwarddiff` after `using ForwardDiff`" - ) + ArgumentError( + "no analytical Jacobian for state modifier $(unwrapped); " * + "use `backend=:forwarddiff` after `using ForwardDiff`" ) - end + ) end return nothing end -function __state_modifiers_jacobian(::Tuple{}, x::AbstractVector) - n = length(x) - return Matrix{eltype(x)}(I, n, n) -end +__state_modifiers_jacobian(::Tuple{}, x::AbstractVector) = + Matrix{eltype(x)}(I, length(x), length(x)) function __state_modifiers_jacobian(modifiers::Tuple, x::AbstractVector) features = x @@ -329,21 +273,16 @@ function __state_modifiers_jacobian(modifiers::Tuple, x::AbstractVector) end function __state_modifier_jacobian(::typeof(NLAT1), x::AbstractVector) - T = eltype(x) - n = length(x) - M = Matrix{T}(I, n, n) + M = Matrix{eltype(x)}(I, length(x), length(x)) @inbounds for i in eachindex(x) - if isodd(i) - M[i, i] = 2 * x[i] - end + isodd(i) && (M[i, i] = 2 * x[i]) end return M end function __state_modifier_jacobian(::typeof(NLAT2), x::AbstractVector) T = eltype(x) - n = length(x) - M = Matrix{T}(I, n, n) + M = Matrix{T}(I, length(x), length(x)) first_i = firstindex(x) @inbounds for i in eachindex(x) if i > first_i && isodd(i) @@ -357,10 +296,8 @@ end function __state_modifier_jacobian(::typeof(NLAT3), x::AbstractVector) T = eltype(x) - n = length(x) - M = Matrix{T}(I, n, n) - first_i = firstindex(x) - last_i = lastindex(x) + M = Matrix{T}(I, length(x), length(x)) + first_i, last_i = firstindex(x), lastindex(x) @inbounds for i in eachindex(x) if first_i < i < last_i && isodd(i) M[i, i] = zero(T) @@ -372,20 +309,13 @@ function __state_modifier_jacobian(::typeof(NLAT3), x::AbstractVector) end function __state_modifier_jacobian(::Pad, x::AbstractVector) - T = eltype(x) n = length(x) - M = zeros(T, n + 1, n) - @inbounds for i in 1:n - M[i, i] = one(T) - end - return M + return vcat(Matrix{eltype(x)}(I, n, n), zeros(eltype(x), 1, n)) end function __state_modifier_jacobian(partial_square::PartialSquare, x::AbstractVector) - T = eltype(x) - n = length(x) - M = Matrix{T}(I, n, n) - threshold = floor(Int, partial_square.eta * n) + M = Matrix{eltype(x)}(I, length(x), length(x)) + threshold = floor(Int, partial_square.eta * length(x)) @inbounds for i in 1:threshold M[i, i] = 2 * x[i] end @@ -393,22 +323,15 @@ function __state_modifier_jacobian(partial_square::PartialSquare, x::AbstractVec end function __state_modifier_jacobian(::typeof(ExtendedSquare), x::AbstractVector) - T = eltype(x) n = length(x) - M = zeros(T, 2n, n) - @inbounds for i in 1:n - M[i, i] = one(T) - M[n + i, i] = 2 * x[i] - end - return M + return vcat(Matrix{eltype(x)}(I, n, n), Matrix(2 .* Diagonal(x))) end function __forwarddiff_closedloop_jacobian!( J::AbstractMatrix, esn::ESN, x::AbstractVector, ps ) ext = Base.get_extension(@__MODULE__, :RCForwardDiffExt) - ext === nothing && error( - "backend=:forwarddiff requires ForwardDiff (`using ForwardDiff`)" - ) + ext === nothing && + error("backend=:forwarddiff requires ForwardDiff (`using ForwardDiff`)") return ext.forwarddiff_closedloop_jacobian!(J, esn, x, ps) end diff --git a/test/Models/jacobian_tests.jl b/test/Models/jacobian_tests.jl index 779febce5..56d376953 100644 --- a/test/Models/jacobian_tests.jl +++ b/test/Models/jacobian_tests.jl @@ -14,8 +14,7 @@ begin n = length(state) J = zeros(eltype(state), n, n) for j in 1:n - state_plus = copy(state) - state_minus = copy(state) + state_plus, state_minus = copy(state), copy(state) state_plus[j] += ε state_minus[j] -= ε forward_plus = first(ReservoirComputing.__closed_loop_step(esn, state_plus, ps)) @@ -26,43 +25,49 @@ begin end function _with_readout_weights(esn, ps, rng) - out_dims = Int(esn.readout.out_dims) - feat_dims = Int(esn.readout.in_dims) + out_dims, feat_dims = Int(esn.readout.out_dims), Int(esn.readout.in_dims) weight = randn(rng, Float32, out_dims, feat_dims) .* 0.05f0 - if haskey(ps.readout, :bias) - return merge(ps, (readout = (weight = weight, bias = zeros(Float32, out_dims)),)) - end - return merge(ps, (readout = (weight = weight,),)) + readout = haskey(ps.readout, :bias) ? + (weight = weight, bias = zeros(Float32, out_dims)) : (weight = weight,) + return merge(ps, (readout = readout,)) + end + + function _check_fd(esn, ps, st, state; atol = _ATOL) + J, st_out = jacobian(esn, state, ps, st) + @test st_out === st + @test size(J) == (length(state), length(state)) + @test J ≈ _finite_diff_jacobian(esn, state, ps) atol = atol + return J end @testset "analytical matches finite differences" begin rng = MersenneTwister(42) esn = ESN( 3, 6, 3; - init_reservoir = scaled_rand, - use_bias = true, - leak_coefficient = 0.7f0, + init_reservoir = scaled_rand, use_bias = true, leak_coefficient = 0.7f0, ) ps, st = setup(rng, esn) ps = _with_readout_weights(esn, ps, rng) state = randn(rng, Float32, 6) - - J_analytical, st_out = jacobian(esn, state, ps, st) - @test st_out === st - @test size(J_analytical) == (6, 6) - @test eltype(J_analytical) === Float32 - @test J_analytical ≈ _finite_diff_jacobian(esn, state, ps) atol = _ATOL + J = _check_fd(esn, ps, st, state) + @test eltype(J) === Float32 + Jbuf = similar(J) + @test first(jacobian!(Jbuf, esn, state, ps, st)) === Jbuf + @test Jbuf ≈ J end - @testset "ForwardDiff backend" begin + @testset "ForwardDiff and tanh_fast" begin rng = MersenneTwister(7) - esn = ESN(2, 5, 2; init_reservoir = scaled_rand) - ps, st = setup(rng, esn) - ps = _with_readout_weights(esn, ps, rng) - state = randn(rng, Float32, 5) - J_analytical, _ = jacobian(esn, state, ps, st) - J_ad, _ = jacobian(esn, state, ps, st; backend = :forwarddiff) - @test J_analytical ≈ J_ad atol = 1.0f-5 + for (activation, seed) in ((tanh, 7), (tanh_fast, 19)) + Random.seed!(rng, seed) + esn = ESN(2, 5, 2, activation; init_reservoir = scaled_rand) + ps, st = setup(rng, esn) + ps = _with_readout_weights(esn, ps, rng) + state = randn(rng, Float32, 5) + J_an, _ = jacobian(esn, state, ps, st) + J_ad, _ = jacobian(esn, state, ps, st; backend = :forwarddiff) + @test J_an ≈ J_ad atol = 1.0f-5 + end end @testset "vector leak" begin @@ -71,22 +76,15 @@ begin esn = ESN(2, 5, 2; init_reservoir = scaled_rand, leak_coefficient = leak) ps, st = setup(rng, esn) ps = _with_readout_weights(esn, ps, rng) - state = randn(rng, Float32, 5) - J_analytical, _ = jacobian(esn, state, ps, st) - @test J_analytical ≈ _finite_diff_jacobian(esn, state, ps) atol = _ATOL + _check_fd(esn, ps, st, randn(rng, Float32, 5)) end @testset "state modifiers" begin rng = MersenneTwister(13) for modifier in (NLAT1, NLAT2, NLAT3, PartialSquare(0.5), ExtendedSquare, Pad(1.0f0)) res_dims = modifier isa Pad ? 5 : 6 - readout_in_dims = if modifier === ExtendedSquare - 12 - elseif modifier isa Pad - 6 - else - res_dims - end + readout_in_dims = modifier === ExtendedSquare ? 12 : + modifier isa Pad ? 6 : res_dims esn = ESN( 2, res_dims, 2; init_reservoir = scaled_rand, @@ -95,32 +93,13 @@ begin ) ps, st = setup(rng, esn) ps = _with_readout_weights(esn, ps, rng) - state = randn(rng, Float32, res_dims) - J_analytical, _ = jacobian(esn, state, ps, st) - @test J_analytical ≈ _finite_diff_jacobian(esn, state, ps) atol = _ATOL + _check_fd(esn, ps, st, randn(rng, Float32, res_dims)) end - end - @testset "NLAT2 sparsity pattern" begin state = Float32[1, 2, 3, 4, 5] M = ReservoirComputing.__state_modifier_jacobian(NLAT2, state) - @test M[3, 3] == 0 - @test M[3, 2] == state[1] - @test M[3, 1] == state[2] - @test M[5, 5] == 0 - @test M[5, 4] == state[3] - @test M[5, 3] == state[4] - end - - @testset "tanh_fast" begin - rng = MersenneTwister(19) - esn = ESN(2, 5, 2, tanh_fast; init_reservoir = scaled_rand) - ps, st = setup(rng, esn) - ps = _with_readout_weights(esn, ps, rng) - state = randn(rng, Float32, 5) - J_analytical, _ = jacobian(esn, state, ps, st) - J_ad, _ = jacobian(esn, state, ps, st; backend = :forwarddiff) - @test J_analytical ≈ J_ad atol = 1.0f-5 + @test M[3, 3] == 0 && M[3, 2] == state[1] && M[3, 1] == state[2] + @test M[5, 5] == 0 && M[5, 4] == state[3] && M[5, 3] == state[4] end @testset "jacobians matches predict" begin @@ -129,21 +108,16 @@ begin ps, st = setup(rng, esn) ps = _with_readout_weights(esn, ps, rng) initialdata = Float32[0.1, -0.2, 0.3] - Js, outputs, st_final = jacobians(esn, 4, ps, st; initialdata) predicted, st_pred = predict(esn, 4, ps, st; initialdata) @test outputs ≈ predicted @test size(Js) == (6, 6, 4) @test st_final.reservoir.carry == st_pred.reservoir.carry - _, st_step = apply(esn, initialdata, ps, st) for t in 1:4 state = ReservoirComputing.__carry_state_vector(st_step) - Jt, _ = jacobian(esn, state, ps, st_step) - @test Js[:, :, t] ≈ Jt atol = 1.0f-6 - if t < 4 - _, st_step = apply(esn, outputs[:, t], ps, st_step) - end + @test Js[:, :, t] ≈ first(jacobian(esn, state, ps, st_step)) atol = 1.0f-6 + t < 4 && ((_, st_step) = apply(esn, outputs[:, t], ps, st_step)) end end @@ -169,17 +143,4 @@ begin @test_throws DimensionMismatch jacobian!(zeros(Float32, 4, 4), esn, state, ps, st) @test_throws ArgumentError jacobians(esn, 0, ps, st; initialdata = Float32[1, 2, 3]) end - - @testset "jacobian!" begin - rng = MersenneTwister(31) - esn = ESN(2, 4, 2; init_reservoir = scaled_rand) - ps, st = setup(rng, esn) - ps = _with_readout_weights(esn, ps, rng) - state = randn(rng, Float32, 4) - J = zeros(Float32, 4, 4) - J_out, _ = jacobian!(J, esn, state, ps, st) - @test J_out === J - J_ref, _ = jacobian(esn, state, ps, st) - @test J ≈ J_ref - end end