optax.projections.projection_halfspace

optax.projections.projection_halfspace(x: Any, a: Any, b: jax.typing.ArrayLike) → Any

Projection onto a halfspace.

Projects a tree x onto the halfspace defined by a tree a and scalar b.

\[\operatorname{argmin}_y \|x - y\|_2^2 \quad \text{subject to} \quad \langle a, y \rangle \leq b\]

Parameters:

• x – tree to project.

• a – tree defining halfspace onto which to project. Must have the same structure as x.

• b – scalar defining halfspace onto which to project.

Returns:

tree with the same structure as x.