Skip to content

Solver refinement#6

Merged
mapi1 merged 9 commits into
mainfrom
solver_refinement
Jun 25, 2026
Merged

Solver refinement#6
mapi1 merged 9 commits into
mainfrom
solver_refinement

Conversation

@mapi1

@mapi1 mapi1 commented Jun 25, 2026

Copy link
Copy Markdown
Collaborator

Adds two NativeSolver knobs and a per-call reduce path for long/large network runs, plus analysis helpers.

New solver kwargs

  • grad_horizon: truncated-BPTT gradient horizon. Severs the carry gradient every W steps so the backward pass stops amplifying past one window (the exp(Tλ) blow-up). Forward is bit-identical.
  • block_size: unified block unit for rematerialization (checkpointing), per-block fold_in streaming noise, and the reduce grain. Renamed from checkpoint_every. Streaming noise keeps the forward noise tensor at one block instead of O(n_steps).
  • reduce= (per-call on solve/prepare): block-level (init, update, finalize) folds that stream a statistic instead of stacking the trajectory, so forward output memory is O(block_size).

Also

  • New adiabatic_scan analysis (bifurcation-diagram sweep, expressed as a jitted lax.scan) and refined Lyapunov helpers.
  • Noise is now part of the scanned-over axis; solver cleanup/housekeeping.
  • Version bump to 0.3.1.

mapi1 added 9 commits June 5, 2026 15:45
…lations:

- tbptt_window: truncated BPTT gradient horizon. Severs the carry gradient
  every W steps so the backward stops amplifying past one window (the
  exp(T*Lambda) blow-up). Forward is bit-identical.
- block_size: the unified block unit for rematerialization (checkpointing),
  per-block fold_in streaming noise, and the reduce grain. Renamed from
  checkpoint_every. Streaming noise reseeds the SDE realization and keeps the
  forward noise tensor at one block instead of O(n_steps).
- reduce= (per-call kwarg on solve/prepare): block-level (init, update,
  finalize) folds that stream a statistic instead of stacking the trajectory,
  so forward output memory is O(block_size).
@mapi1 mapi1 merged commit 41085a5 into main Jun 25, 2026
7 of 9 checks passed
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.

1 participant