optax.projections.projection_box

optax.projections.projection_box(tree: Any, lower: Any, upper: Any) → Any

Projection onto box constraints.

\[\underset{p}{\text{argmin}} ~ \|x - p\|_2^2 \quad \textrm{subject to} \quad \text{lower} \le p \le \text{upper}\]

where \(x\) is the input tree.

Parameters:

• tree – tree to project.

• lower – lower bound, a scalar or tree with the same structure as tree.

• upper – upper bound, a scalar or tree with the same structure as tree.

Returns:

projected tree, with the same structure as tree.