Skip to content

feat: closed-loop ESN Jacobian API (#169) - #540

Open
Saswatsusmoy wants to merge 2 commits into
SciML:masterfrom
Saswatsusmoy:Saswatsusmoy/169-jacobian-trajectory
Open

Saswatsusmoy wants to merge 2 commits into
SciML:masterfrom
Saswatsusmoy:Saswatsusmoy/169-jacobian-trajectory

Conversation

@Saswatsusmoy

Copy link
Copy Markdown
Contributor

Checklist

  • Appropriate tests were added
  • Any code changes were done in a way that does not break public API
  • All documentation related to code changes were updated
  • The new code follows the contributor guidelines, in particular the SciML Style Guide and COLPRAC
  • Any new documentation only uses public API

Additional context

Adds jacobian / jacobian! / jacobians for the closed-loop reservoir map of a discrete ESN (in_dims == out_dims), so Jacobians can be taken along generative trajectories for Lyapunov analysis (Pathak2017).

  • Analytical default for standard activations and modifiers ((), NLAT1/2/3, Pad, PartialSquare, ExtendedSquare)
  • Optional backend=:forwarddiff via RCForwardDiffExt (ForwardDiff weakdep)
  • jacobians uses the same autoregressive feedback as predict
  • Deep / Hybrid / Continuous / Extend are out of scope for this PR
J, st = jacobian(esn, state, ps, st)
Js, outputs, st = jacobians(esn, steps, ps, st; initialdata)

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

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread src/jacobian.jl

T = eltype(x)
leak = __format_leak(T, cell.leak_coefficient)
x_new = __one_minus_leak(T, leak) .* x .+ leak .* cell.activation.(preactivation)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Comment thread src/jacobian.jl
Comment on lines +389 to +390
@inbounds for i in 1:threshold
M[i, i] = 2 * x[i]

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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.
@MartinuzziFrancesco

Copy link
Copy Markdown
Collaborator

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!

@Saswatsusmoy

Copy link
Copy Markdown
Contributor Author

Agreed — I'll convert this PR to a documentation example instead of adding Jacobian APIs to the main library.

Plan:

  • drop the core jacobian / jacobian! / jacobians API and the ForwardDiff extension
  • add a docs example (closed-loop ESN trajectory → Jacobian / Lyapunov-style analysis with existing Lux + predict APIs)
  • keep #169 in mind as docs-only for now; a utilities package can wait until that direction is clearer

I'll push the reworked branch shortly.

@ChrisRackauckas

Copy link
Copy Markdown
Member

🤖 Automated review from an AI agent running as @ChrisRackauckas — not written or reviewed by Chris. It is posted so that you can act on it. Chris or a maintainer may disagree.

Verdict: changes needed before merge (reviewed at 752c311).

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

  • Risk: medium
  • Blast radius: Three new exported APIs and a ForwardDiff extension for discrete ESNs; incorrect or unsafe modifier handling affects Jacobian/Lyapunov analysis, and missing documentation blocks the documentation build. Existing prediction code is not changed.
  • Evidence: Read the full diff, linked issue, all conversation comments, inline comments, and review. Locally ran the existing Models group on Julia 1.12.7: 653/653 assertions passed across 12 files, including 43/43 in Models/jacobian_tests.jl. Runic returned exit 0 for all six changed Julia files. Additional probes confirmed the missing jacobian! docstring, the undeclared-public ForwardDiff name, and the PartialSquare bounds error. CI comparisons below distinguish introduced failures from historical base failures.
  • Merge: needs changes (complete the maintainer-requested scope rework; if retaining the core APIs, address findings 2–5).

Local commands and observed output:

TMPDIR=$HOME/tmp JULIA_NUM_THREADS=4 JULIA_NUM_PRECOMPILE_TASKS=2 JULIA_PKG_PRECOMPILE_AUTO=0 GROUP=Models timeout 3600 /home/crackauc/.juliaup/bin/julia +1.12 --project=repo -e 'using Pkg; Pkg.test()'
Models/jacobian_tests.jl | 43  43  1m23.3s
Testing ReservoirComputing tests passed
# Aggregate of the 12 printed file summaries: 653 passed / 653 total; exit 0.

/home/crackauc/.juliaup/bin/julia +1.12 --project=/home/crackauc/.julia/environments/runic -e 'using Runic; exit(Runic.main(vcat(["--check", "--diff"], ARGS)))' repo/docs/pages.jl repo/ext/RCForwardDiffExt.jl repo/src/ReservoirComputing.jl repo/src/jacobian.jl repo/test/Models/jacobian_tests.jl repo/test/qa/qa.jl
# No output; exit 0.

/home/crackauc/.juliaup/bin/julia +1.12 --check-bounds=yes --project=probe-env edge-probes.jl
vector leak: state size=(5, 1); AD size=(5, 5); max error=1.4901161e-8
jacobian! docstring registered: false
ForwardDiff.jacobian public: false
PartialSquare(2) apply size: (2, 1)
PartialSquare(2) jacobian error: BoundsError: attempt to access 5-element Vector{Float32} at index [6]

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 4732332 run https://github.com/SciML/ReservoirComputing.jl/actions/runs/34766888897:

Failing head check(s) Base comparison and cause
Documentation Head fails with no docs found for 'jacobian!' and makedocs ... [:docs_block]; base Documentation passed. Introduced by this PR.
tests / QA (julia 1, ubuntu-latest) / Tests - QA (Julia 1) Base passed; head reports isempty([:jacobian!]) and NonPublicQualifiedAccessException for ForwardDiff.jacobian in the new extension. Introduced by this PR.
tests / Interface (julia 1, ubuntu-latest), tests / Interface (julia lts, ubuntu-latest) Both fail on head and base with the same ridge-regression QR solver error.
tests / Interface (julia 1, macos-latest), tests / Interface (julia lts, macos-latest) Both fail on head and base with the same ridge-regression QR solver error.
tests / Interface (julia 1, windows-latest), tests / Interface (julia lts, windows-latest) Both fail on head and base with the same ridge-regression QR solver error.
tests / Core (julia 1, windows-latest), tests / Core (julia lts, windows-latest) Both fail on head and base with the same ridge-regression QR solver error.
tests / Extensions (julia 1, windows-latest), tests / Extensions (julia lts, windows-latest) Both fail on head and base with the same ridge-regression QR solver error.
tests / Layers (julia 1, windows-latest), tests / Layers (julia lts, windows-latest) Both fail on head and base with the same ridge-regression QR solver error.
tests / Models (julia 1, windows-latest), tests / Models (julia lts, windows-latest) Both fail on head and base with the same ridge-regression QR solver error.

The 14 historical failures report ArgumentError: solver LinearSolve.QRFactorization{LinearAlgebra.NoPivot} failed to solve the ridge regression system in test/Interface/array_workflow_tests.jl:111; they are not Jacobian regressions. Current master a02f2c5 has merged the QR fallback fix and passes all corresponding test checks, QA, and Documentation in the linked current-master runs. Head Runic, spelling, and downgrade checks passed.

Findings

  1. Scope remains contrary to the maintainer's requested rework — src/ReservoirComputing.jl:113. Maintainer MartinuzziFrancesco requested a documentation example rather than core APIs; the author explicitly agreed to remove these exports and the ForwardDiff extension. This head still adds all of them and contains only an API reference page, not the promised example. Complete that rework before merging; the PR title/body should describe the resulting documentation example.
  2. Missing docstring blocks Documentation and QA — src/jacobian.jl:46. The preceding docstring is attached only to jacobian; mentioning jacobian! in its signature block does not document that separate binding. Both CI errors and the local documentation metadata probe confirm this. If retained, attach a docstring to jacobian! and keep its rendered @docs entry.
  3. PartialSquare derivative indexes beyond the state — src/jacobian.jl:319. PartialSquare(2) is accepted by existing evaluation and squares every component, but the derivative iterates to floor(eta*n) without bounding it by n. For a five-component state, the local checked run throws at index 6; the @inbounds annotation makes this unsafe in ordinary unchecked execution. Iterate over valid indices subject to the threshold, or clamp the threshold, and add an edge-case test matching existing forward behavior.
  4. New extension fails the public-API gate — ext/RCForwardDiffExt.jl:13. ForwardDiff.jacobian is documented upstream but is not exported or declared public in the resolved package (also false in local ForwardDiff 1.4.6). The repository's existing ExplicitImports check rejects this access. If retaining the extension, obtain the upstream public declaration and an appropriate compatibility floor instead of silently bypassing QA; the agreed docs-only conversion also removes this extension failure.
  5. New public API has no corresponding minor version bump — Project.toml:3. The PR exports three new names while leaving the version at 0.12.42. Under the requested review policy, retaining these APIs requires a minor release bump; converting to the agreed documentation-only scope avoids adding this public surface.

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
Conversation: local Claude Code session 3cd6500a-1f81-46b5-ac0b-c466e15b6a53 on Chris's Mac (session ID, no URL)

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Feature Request: Jacobian along Trajectory

3 participants