feat: closed-loop ESN Jacobian API (#169) - #540
Saswatsusmoy wants to merge 2 commits into
Conversation
Add jacobian/jacobian!/jacobians for the autonomous reservoir map of a trained ESN (analytical default; ForwardDiff weakdep fallback), with Models tests and API docs.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9f1575f7f7
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| T = eltype(x) | ||
| leak = __format_leak(T, cell.leak_coefficient) | ||
| x_new = __one_minus_leak(T, leak) .* x .+ leak .* cell.activation.(preactivation) |
There was a problem hiding this comment.
Keep heterogeneous-leak closed-loop states vector-shaped
When leak_coefficient is a vector, __format_leak reshapes it to n×1, so broadcasting it with the vector x and vector preactivation produces an n×n x_new. Consequently backend=:forwarddiff differentiates a length-n² map and cannot assign its result into the required n×n Jacobian; the newly added vector-leak finite-difference test also reaches this malformed map. Flatten the leak for this vector-only closed-loop path (or otherwise keep the operands one-dimensional).
Useful? React with 👍 / 👎.
| @inbounds for i in 1:threshold | ||
| M[i, i] = 2 * x[i] |
There was a problem hiding this comment.
Bound the PartialSquare derivative to the state length
For PartialSquare(eta) with eta > 1, the existing forward modifier simply squares every state component, but this loop uses floor(Int, eta * n) without capping it and writes past the n×n matrix. Such a model can run normally in prediction yet jacobian(...; backend=:analytical) throws a BoundsError; cap the derivative loop at n (or validate eta consistently in the modifier constructor).
Useful? React with 👍 / 👎.
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.
|
This is a nice direction, but I don't think it should be in main library. Perhaps it is time to think about a ReservoirComputingUtilities library However, I do think it would make for a fantastic example in the documentation! |
|
Agreed — I'll convert this PR to a documentation example instead of adding Jacobian APIs to the main library. Plan:
I'll push the reworked branch shortly. |
Verdict: changes needed before merge (reviewed at Adds analytical and optional ForwardDiff closed-loop ESN Jacobians, trajectory Jacobians, and API documentation. The full Models group passes locally, but the agreed documentation-only rework is absent, Documentation and QA regress, and PartialSquare has an unchecked bounds defect. Risk assessment
Local commands and observed output: The initial Models attempt received SIGTERM during dependency precompilation; the retry above completed successfully. A separate execution of the unchanged new Jacobian test file also passed all 43 assertions. The earlier bot comment claiming heterogeneous leak produces an n×n state was not reproduced: the state is n×1 and both differentiation paths succeed. No numerical speedup claim is made by this PR. ForwardDiff and ReservoirComputing have MIT licenses; no new copyleft dependency or weakened existing test was found. Full local QA, docs, GPU, and other test groups were not run; the QA/docs conclusions use their actual CI error logs. This is a feature addition, so there is no bug-fix failing-before/passing-after claim. CI comparison: head test run https://github.com/SciML/ReservoirComputing.jl/actions/runs/34939432186 versus the merge-base
The 14 historical failures report Findings
Push a fix and the PR is reviewed again automatically at the new head. 🤖 Posted by an AI agent — harness: Claude Code · model: claude-opus-5-5[1m] (fleet master); review by Codex CLI 0.157.1 / gpt-6-astra |
Checklist
Additional context
Adds
jacobian/jacobian!/jacobiansfor the closed-loop reservoir map of a discreteESN(in_dims == out_dims), so Jacobians can be taken along generative trajectories for Lyapunov analysis (Pathak2017).(), NLAT1/2/3,Pad,PartialSquare,ExtendedSquare)backend=:forwarddiffviaRCForwardDiffExt(ForwardDiff weakdep)jacobiansuses the same autoregressive feedback aspredictExtendare out of scope for this PRCloses #169