optax.projections.projection_hyperplane

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

Projection onto a hyperplane.

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

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

Parameters:

• x – tree to project.

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

• b – scalar defining hyperplane onto which to project.

Returns:

tree with the same structure as x.