Jaxopt

New features

- Pjit-based example of data parallel training using Flax, by Felipe Llinares.

Bug fixes and enhancements

- Support for GPU and state of the art adversarial training algorithm (PGD) on the robust_training.py example, by Fabian Pedregosa.
- Update line search in LBFGS to use jit and unroll from LBFGS, by Ian Williamson.
- Support dynamic maximum iteration count in iterative solvers, by Roy Frostig.
- Fix tree_where for singleton pytrees, by Louis Béthune.
- Remove QuadraticProg in projections and set ``init_params=None`` by default in QP solvers, by Louis Béthune.
- Add missing 'value' attribute in LbfgsState, by Mathieu Blondel.

Contributors

Felipe Llinares, Fabian Pedregosa, Ian Williamson, Louis Bétune, Mathieu Blondel, Roy Frostig.

New features

- jaxopt.LBFGS
- jaxopt.BacktrackingLineSearch
- jaxopt.GaussNewton
- jaxopt.NonlinearCG

Bug fixes and enhancements

- Support implicit AD in higher-order differentiation.

Contributors

Amir Saadat, Fabian Pedregosa, Geoffrey Négiar, Hyunsung Lee, Mathieu Blondel, Roy Frostig.

New features

- Quadratic programming solvers jaxopt.CvxpyQP, jaxopt.OSQP, jaxopt.BoxOSQP and jaxopt.EqualityConstrainedQP
- Iterative refinement

New examples

- Resnet example with Flax and JAXopt.

Bug fixes and enhancements

- Prevent recompilation of loops in solver.run if executing without jit.
- Prevents recomputation of gradient in OptaxSolver.
- Make solver.update jittable and ensure output states are consistent.
- Allow Callable for the stepsize argument in jaxopt.ProximalGradient, jaxopt.ProjectedGradient and jaxopt.GradientDescent.

Deprecated features

- jaxopt.QuadraticProgramming is deprecated and will be removed in v0.3. Use jaxopt.CvxpyQP, jaxopt.OSQP, jaxopt.BoxOSQP and jaxopt.EqualityConstrainedQP instead.

Contributors

Fabian Pedregosa, Felipe Llinares, Geoffrey Negiar, Louis Bethune, Mathieu Blondel, Vikas Sindhwani.

New features

- Added solver jaxopt.ArmijoSGD
- Added example Deep Equilibrium (DEQ) model in Flax with Anderson acceleration.
- Added example Comparison of different SGD algorithms.

Bug fixes

- Allow non-jittable proximity operators in jaxopt.ProximalGradient
- Raise an exception if a quadratic program is infeasible or unbounded

Contributors

Fabian Pedregosa, Louis Bethune, Mathieu Blondel.

Classes

- jaxopt.AndersonAcceleration
- jaxopt.AndersonWrapper
- jaxopt.Bisection
- jaxopt.BlockCoordinateDescent
- jaxopt.FixedPointIteration
- jaxopt.GradientDescent
- jaxopt.MirrorDescent
- jaxopt.OptaxSolver
- jaxopt.PolyakSGD
- jaxopt.ProjectedGradient
- jaxopt.ProximalGradient
- jaxopt.QuadraticProgramming
- jaxopt.ScipyBoundedLeastSquares
- jaxopt.ScipyBoundedMinimize
- jaxopt.ScipyLeastSquares
- jaxopt.ScipyMinimize
- jaxopt.ScipyRootFinding
- Implicit differentiation

Examples

- Binary kernel SVM with intercept.
- Image classification example with Flax and JAXopt.
- Image classification example with Haiku and JAXopt.
- VAE example with Haiku and JAXopt.
- Implicit differentiation of lasso.
- Multiclass linear SVM (without intercept).
- Non-negative matrix factorizaton (NMF) using alternating minimization.
- Dataset distillation.
- Implicit differentiation of ridge regression.
- Robust training.
- Anderson acceleration of gradient descent.
- Anderson acceleration of block coordinate descent.
- Anderson acceleration in application to Picard–Lindelöf theorem.

Contributors

Fabian Pedregosa, Felipe Llinares, Robert Gower, Louis Bethune, Marco Cuturi, Mathieu Blondel, Peter Hawkins, Quentin Berthet, Roy Frostig, Ta-Chu Kao