jaxopt.EqualityConstrainedQP

class jaxopt.EqualityConstrainedQP(matvec_Q=None, matvec_A=None, fun=None, solve=<function solve_gmres>, refine_regularization=0.0, refine_maxiter=10, maxiter=1000, tol=1e-05, implicit_diff_solve=None, jit=True)[source]

Quadratic programming with equality constraints only.

Supports implicit differentiation, matvec, pytrees, and quadratic functions. Can benefit from GPU/TPU acceleration.

If the algorithm diverges on some instances, it might be useful to tweak the refine_regularization parameter.

Parameters

  • matvec_Q (Optional[Callable]) –

  • matvec_A (Optional[Callable]) –

  • fun (Optional[Callable]) –

  • solve (Callable) –

  • refine_regularization (float) –

  • refine_maxiter (int) –

  • maxiter (int) –

  • tol (float) –

  • implicit_diff_solve (Optional[Callable]) –

  • jit (bool) –

matvec_Q

a Callable matvec_Q(params_Q, u). By default, matvec_Q(Q, u) = dot(Q, u), where Q = params_Q. matvec_Q incompatible with the specification of fun. The shape of primal variables may be inferred from params_obj = (matvec_Q, c).

Type

Optional[Callable]

matvec_A

a Callable matvec_A(params_A, u). By default, matvec_A(A, u) = dot(A, u), where A = params_A.

Type

Optional[Callable]

fun

(optional) a function with signature fun(params, params_obj) that is promised to be a quadratic polynomial convex with respect to params, i.e fun can be written

fun(x, params_obj) = 0.5*jnp.dot(x, jnp.dot(Q, x)) + jnp.dot(c, x) + cste

with params_obj a pytree that contains the parameters of the objective function. (Q, c) do not need to be explicited in params_obj by the user: c will be inferred by Jaxopt,

and the operator x -> Qx will be computed upon request.

fun incompatible with the specification of matvec_Q. Note that the shape of primal cannot be inferred from params_obj anymore, so the user should provide it in init_params.

Type

Optional[Callable]

solve

a Callable to solve linear systems, that accepts matvecs (default: linear_solve.solve_gmres).

Type

Callable

Parameters

  • matvec (Callable) –

  • b (Any) –

  • ridge (Optional[float]) –

  • init (Optional[Any]) –

  • tol (float) –

Return type

Any

refine_regularization

a float (default: 0.) used to regularize the system. Useful for badly conditioned problems that lead to divergence. IterativeRefinement is used to correct the error introduced.

Type

float

refine_maxiter

maximum number of refinement steps, when refine_regularization is not 0.

Type

int

maxiter

maximum number of iterations.

Type

int

tol

tolerance (stoping criterion).

Type

float

implicit_diff

whether to enable implicit diff or autodiff of unrolled iterations.

implicit_diff_solve

the linear system solver to use.

Type

Optional[Callable]

jit

whether to JIT-compile the optimization loop (default: “auto”).

Type

bool

__init__(matvec_Q=None, matvec_A=None, fun=None, solve=<function solve_gmres>, refine_regularization=0.0, refine_maxiter=10, maxiter=1000, tol=1e-05, implicit_diff_solve=None, jit=True)
Parameters

  • matvec_Q (Optional[Callable]) –

  • matvec_A (Optional[Callable]) –

  • fun (Optional[Callable]) –

  • solve (Callable) –

  • refine_regularization (float) –

  • refine_maxiter (int) –

  • maxiter (int) –

  • tol (float) –

  • implicit_diff_solve (Optional[Callable]) –

  • jit (bool) –

Return type

None

Methods

__init__([matvec_Q, matvec_A, fun, solve, ...])

attribute_names()

attribute_values()

l2_optimality_error(params, params_obj, ...)

Computes the L2 norm of the KKT residuals.

run([init_params, params_obj, params_eq])

Solves 0.5 * x^T Q x + c^T x subject to Ax = b.

solve(b[, ridge, init, tol])

Solves A x = b using gmres.

Attributes

fun

implicit_diff_solve

jit

matvec_A

matvec_Q

maxiter

refine_maxiter

refine_regularization

tol

l2_optimality_error(params, params_obj, params_eq)[source]

Computes the L2 norm of the KKT residuals.

Parameters

  • params (KKTSolution) –

  • params_obj (Optional[Any]) –

  • params_eq (Optional[Any]) –

run(init_params=None, params_obj=None, params_eq=None)[source]

Solves 0.5 * x^T Q x + c^T x subject to Ax = b.

This solver returns both the primal solution (x) and the dual solution.

Parameters

  • init_params (Optional[KKTSolution]) – (optional), used to infer the shape of primal variables. Mandatory if fun is not None, since the shape of primal variables cannot be inferred from params_obj

  • params_obj (Optional[Any]) – parameters of the quadratic objective, can be: a tuple (Q, c) with Q a pytree of matrices, or a tuple (params_Q, c) if matvec_Q is provided, or an arbitrary pytree if fun is provided.

  • params_eq (Optional[Any]) – parameters of the equality constraints, can be: a tuple (A, b) with A a pytree of matrices, or a tuple (params_A, b) if matvec_A is provided.

Return type

OptStep

Returns

(params, state), where params = (primal_var, dual_var_eq, None)

solve(b, ridge=None, init=None, tol=1e-05, **kwargs)

Solves A x = b using gmres.

Parameters

  • matvec (Callable) – product between A and a vector.

  • b (Any) – pytree.

  • ridge (Optional[float]) – optional ridge regularization.

  • init (Optional[Any]) – optional initialization to be used by gmres.

  • **kwargs – additional keyword arguments for solver.

  • tol (float) –

Return type

Any

Returns

pytree with same structure as b.