optax.projections.projection_linf_ball

optax.projections.projection_linf_ball#

optax.projections.projection_linf_ball(
tree: Any,
scale: jax.typing.ArrayLike = 1,
) → Any[source]#

Projection onto the l-infinity ball.

This function solves the following constrained optimization problem, where x is the input tree.

\[\underset{y}{\text{argmin}} ~ \|x - y\|_2^2 \quad \textrm{subject to} \quad \|y\|_{\infty} \le \text{scale}\]
Parameters:

  • tree – tree to project.

  • scale – radius of the ball.

Returns:

projected tree, with the same structure as tree.