Repository navigation

Add an L-BFGS rule with a ring-buffered memory and an mlx dispatch - #162

Merged

jessegrabowski merged 31 commits into

pymc-devs:mainfrom

jessegrabowski:lbfgs-direction-rule

Oct 2, 2026

Merged

Add an L-BFGS rule with a ring-buffered memory and an mlx dispatch#162
jessegrabowski merged 31 commits into
pymc-devs:mainfrom
jessegrabowski:lbfgs-direction-rule

Commits on Sep 19, 2026
- Let state_for stack a history axis
  jessegrabowski
  committed
  f8f3ad3
- Add an LBFGSDirection op with a scan inner graph
  jessegrabowski
  committed
  2d8e5b7
- Take the recursion's dot products through BLAS
  jessegrabowski
  committed
  e4ebd9f
- Dispatch LBFGSDirection to mlx as a loop over ring rows
  jessegrabowski
  committed
  e4893cf
- Test the direction op against the dense BFGS matrix
  jessegrabowski
  committed
  63c5d6f
- Keep each parameter's dtype through the recursion
  jessegrabowski
  committed
  959b004
- Compute the curvatures with the dot the rule's guard uses
  jessegrabowski
  committed
  75af337
- Allocate optimizer state at the parameter's declared dtype
  jessegrabowski
  committed
  5e86542
- Add lbfgs_updates with a ring memory and curvature guard
  jessegrabowski
  committed
  fceed2a
- Raise the mlx floor to 0.32.1 and dot through its vector matmul
  jessegrabowski
  committed
  8b5a8de
- Set the mlx floor to 0.32.2
  jessegrabowski
  committed
  13eba45
- Add the lbfgs alias
  jessegrabowski
  committed
  05d2ffb
- Build literal count and gamma at their dtype instead of casting
  jessegrabowski
  committed
  a20f465
- Require pytensor 3.3.2
  jessegrabowski
  committed
  914fffd
Commits on Oct 2, 2026
- Cast the mlx recursion to each parameter's dtype
  jessegrabowski
  committed
  41274c2
- Cache each admitted pair's curvature in rule state
  jessegrabowski
  committed
  5177de8
- Require pytensor 3.3.3
  jessegrabowski
  committed
  215e6ba
- Move the L-BFGS test references into their own module
  jessegrabowski
  committed
  08d6215
- Require static slot counts in LBFGSDirection
  jessegrabowski
  committed
  c6d8338
- Build gamma at the curvature dtype
  jessegrabowski
  committed
  f443bc5
- Carry the L-BFGS identity scale in rule state
  jessegrabowski
  committed
  7744434
- Refuse L-BFGS pairs whose curvature overflows
  jessegrabowski
  committed
  2a29785
- Test L-BFGS on a one-slot ring
  jessegrabowski
  committed
  20e231d
- Test that L-BFGS rejects a negligible positive curvature
  jessegrabowski
  committed
  3e036f3
- Test L-BFGS with parameters of mixed dtypes
  jessegrabowski
  committed
  e5629c2
- Test that L-BFGS resumes from a checkpoint
  jessegrabowski
  committed
  883e742
- Run lbfgs at its default rate in the loss-reduction test
  jessegrabowski
  committed
  218b635
- Test lbfgs with its step clipped after it
  jessegrabowski
  committed
  e2fbdf3
- Test that skip_if holds back the L-BFGS state
  jessegrabowski
  committed
  a6fb4d1
- Document the fixed-step contract of lbfgs
  jessegrabowski
  committed
  18ee016
- List LBFGSDirection in the optim API docs
  jessegrabowski
  committed
  0bf07b0

Commits

Commits on Sep 19, 2026

Commits on Oct 2, 2026