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..6d829cdea --- /dev/null +++ b/src/jacobian.jl @@ -0,0 +1,337 @@ +@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")) + __require_esn_closedloop_io(esn) + 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 + +__jacobian_state_vector(state::AbstractVector) = state + +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 = 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 + +__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 + +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 __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) + J .= q.input_matrix * __readout_jacobian(esn.readout, q.z, ps.readout, M) + J .+= q.reservoir_matrix # densifies if `reservoir_matrix` is sparse + return __finalize_leak_jacobian!( + J, q.leak, __activation_derivative(esn.reservoir.cell.activation, q.preactivation) + ) +end + +function __finalize_leak_jacobian!(J::AbstractMatrix, leak::Number, dφ::AbstractVector) + J .*= leak .* dφ + @inbounds for i in axes(J, 1) + J[i, i] += one(eltype(J)) - leak + end + return J +end + +function __finalize_leak_jacobian!(J::AbstractMatrix, leak::AbstractArray, dφ::AbstractVector) + α = vec(leak) + length(α) == length(dφ) || throw( + DimensionMismatch( + "leak_coefficient length $(length(α)) must match reservoir size $(length(dφ))" + ) + ) + J .*= α .* dφ + @inbounds for i in axes(J, 1) + J[i, i] += one(eltype(J)) - α[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 + 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)) +__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( + 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) + unwrapped isa Extend && throw( + ArgumentError( + "closed-loop jacobian does not support `Extend` " * + "(autonomous state is not the reservoir carry alone)" + ) + ) + hasmethod(__state_modifier_jacobian, Tuple{typeof(unwrapped), AbstractVector}) || + throw( + ArgumentError( + "no analytical Jacobian for state modifier $(unwrapped); " * + "use `backend=:forwarddiff` after `using ForwardDiff`" + ) + ) + end + return nothing +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 + 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) + M = Matrix{eltype(x)}(I, length(x), length(x)) + @inbounds for i in eachindex(x) + isodd(i) && (M[i, i] = 2 * x[i]) + end + return M +end + +function __state_modifier_jacobian(::typeof(NLAT2), x::AbstractVector) + T = eltype(x) + M = Matrix{T}(I, length(x), length(x)) + 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) + 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) + 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) + n = length(x) + return vcat(Matrix{eltype(x)}(I, n, n), zeros(eltype(x), 1, n)) +end + +function __state_modifier_jacobian(partial_square::PartialSquare, x::AbstractVector) + 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 + return M +end + +function __state_modifier_jacobian(::typeof(ExtendedSquare), x::AbstractVector) + n = length(x) + 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`)") + 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..56d376953 --- /dev/null +++ b/test/Models/jacobian_tests.jl @@ -0,0 +1,146 @@ +# 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, state_minus = copy(state), 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, feat_dims = Int(esn.readout.out_dims), Int(esn.readout.in_dims) + weight = randn(rng, Float32, out_dims, feat_dims) .* 0.05f0 + 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, + ) + ps, st = setup(rng, esn) + ps = _with_readout_weights(esn, ps, rng) + state = randn(rng, Float32, 6) + 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 and tanh_fast" begin + rng = MersenneTwister(7) + 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 + 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) + _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 = modifier === ExtendedSquare ? 12 : + modifier isa Pad ? 6 : res_dims + 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) + _check_fd(esn, ps, st, randn(rng, Float32, res_dims)) + end + + state = Float32[1, 2, 3, 4, 5] + M = ReservoirComputing.__state_modifier_jacobian(NLAT2, state) + @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 + 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) + @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 + + @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 +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,