Skip to content

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

jessegrabowski merged 31 commits into
pymc-devs:mainfrom
jessegrabowski:lbfgs-direction-rule

Conversation

@jessegrabowski

@jessegrabowski jessegrabowski commented Sep 17, 2026 •

Copy link
Copy Markdown
Member

Adds L-BFGS at a fixed step size. lbfgs_updates in rules.py keeps per-parameter stacks of the last memory_size parameter and gradient differences as a ring, written in place with set_subtensor, and admits a pair only when y . s > eps * y . y, so the inverse-Hessian estimate stays positive definite. LBFGSDirection, a SymbolicOp in optim/lbfgs.py, applies the two-loop recursion to the gradient with two scans over the ring order and gets a Python-loop dispatch in dispatch/mlx/lbfgs.py, since Scan has no mlx dispatch. The lbfgs alias takes a rate or a schedule like the other rules.

The ring slot is a one-element index vector rather than a scalar. pytensor's mlx Subtensor dispatch cannot trace a scalar index (pymc-devs/pytensor#2422), and the advanced-indexing ops can.

The mlx CI install now requires mlx>=0.32.2. Below that, mlx's vector matmul ran one threadgroup and a 12.6M-element dot took 155 ms instead of 0.4 ms (ml-explore/mlx#3580).

state_for takes history_size for the stacked buffers and scalar_state takes a dtype for the pair count. Tests check the op against the dense BFGS matrix built from its definition and the rule against the secant condition and a quadratic's closed-form minimizer, on numba and on mlx.

The line search is the next PR. Part of #58.


📚 Documentation preview 📚: https://pytensor-ml--162.org.readthedocs.build/en/162/

@jessegrabowski
jessegrabowski merged commit 40ae351 into pymc-devs:main Oct 2, 2026
14 of 22 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