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. |
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