JAXopt is a Python library providing hardware-accelerated, batchable, and differentiable optimizers built on JAX. It offers a wide range of solvers for convex and non-convex optimization problems, suitable for machine learning and scientific computing, including gradient descent, L-BFGS, and quadratic programming. The current version is 0.8.5, and the project maintains a frequent release cadence with bug fixes and new features.
pip install jaxoptVerified import paths — ran on the pinned version, not inferred.
This example demonstrates how to use `GradientDescent` from JAXopt to minimize a simple quadratic loss function. It sets up dummy data, initializes model parameters, and then runs the optimizer to find the optimal parameters. This covers the basic workflow of defining an objective, choosing a solver, and executing it.
Upgrade your Python environment to version 3.10 or later.
If your code or custom JAXopt extensions utilize `jax.pure_callback`, explicitly define the `vmap_method` argument (e.g., `vmap_method='sequential'`) to avoid future breaking changes.
Ensure you are using JAXopt version 0.8.4 or later if you are experiencing issues with PyTree-structured parameters in constrained optimization problems.
Update any code or examples referencing the 'boston' dataset to use alternative datasets, such as those recommended by scikit-learn (e.g., California housing dataset), or other suitable benchmarks.
Ensure that any custom objective function, and especially its outputs, are fully JAX-compatible, using `jax.numpy` and JAX control flow (`jax.lax.cond`, `jax.lax.scan`) instead of Python/NumPy equivalents, or explicitly manage jittability. The output `carry` of the scanned function must have the same type, shape, and dtype as the input `carry`.
Convert scalar Python floats to JAX arrays using `jax.numpy.array()` or `jnp.asarray()` before passing them to JAXopt solvers or JAX-jitted functions. Ensure all initial parameters are JAX arrays.
Provide an initial guess for the primal variables (e.g., `init_params=jnp.zeros(problem_dim)`) to the solver's `run` method when `fun` is used to define the quadratic program, especially if the primal variable shape cannot be inferred otherwise.
Ensure that all possible return paths within JAX's conditional control flow functions (`jax.lax.cond`, `jax.lax.while_loop`) return outputs with identical PyTree structures (same number of leaves, same nested structure), and that corresponding leaves have the same shape and dtype. Padding or explicit type casting might be necessary to unify outputs.