Add BFGS optimizer example and benchmark vs. Jaxopt.

google-research / dex-lang

Research language for array processing in the Haskell/ML family

BSD 3-Clause "New" or "Revised" License

1.56k stars 106 forks source link

Add BFGS optimizer example and benchmark vs. Jaxopt. #1302

Closed emilyfertig closed 1 year ago

emilyfertig commented 1 year ago

Pair programmed with @axch .

Run time of the Dex logistic regression example went from ~240 s to 26 s following 289a43e922666114795cdef8c32d8f2a9026a5a9.

emilyfertig commented 1 year ago

I removed the trivial logistic regression benchmark because it took 7 s to compile. I didn’t use dex_test_mode for the Rosenbrock example since the time isn’t very sensitive to max iterations/max line search iterations.