Projections

Projections can be used to perform constrained optimization. The Euclidean projection onto a set \(\mathcal{C}\) is:

\[\text{proj}_{\mathcal{C}}(u) := \underset{v}{\text{argmin}} ~ \|u - v\|^2_2 \textrm{ subject to } v \in \mathcal{C}.\]

For instance, here is an example how we can project parameters to the non-negative orthant:

>>> import optax
>>> import jax
>>> import jax.numpy as jnp
>>> num_weights = 2
>>> xs = jnp.array([[-1.8, 2.2], [-2.0, 1.2]])
>>> ys = jnp.array([0.5, 0.8])
>>> optimizer = optax.adam(learning_rate=1e-3)
>>> params = {'w': jnp.zeros(num_weights)}
>>> opt_state = optimizer.init(params)
>>> loss = lambda params, x, y: jnp.mean((params['w'].dot(x) - y) ** 2)
>>> grads = jax.grad(loss)(params, xs, ys)
>>> updates, opt_state = optimizer.update(grads, opt_state)
>>> params = optax.apply_updates(params, updates)
>>> params = optax.projections.projection_non_negative(params)

Available projections

projection_box(tree, lower, upper)	Projection onto box constraints.
projection_hypercube(tree[, scale])	Projection onto the (unit) hypercube.
projection_l1_ball(tree[, scale])	Projection onto the l1 ball.
projection_l1_sphere(tree[, scale])	Projection onto the l1 sphere.
projection_l2_ball(tree[, scale])	Projection onto the l2 ball.
projection_l2_sphere(tree[, scale])	Projection onto the l2 sphere.
projection_linf_ball(tree[, scale])	Projection onto the l-infinity ball.
projection_non_negative(tree)	Projection onto the non-negative orthant.
projection_simplex(tree[, scale])	Projection onto a simplex.
projection_vector(x, a)	Projection onto a vector.
projection_hyperplane(x, a, b)	Projection onto a hyperplane.
projection_halfspace(x, a, b)	Projection onto a halfspace.