optax.projections.projection_vector

optax.projections.projection_vector(x: Any, a: Any) → Any

Projection onto a vector.

Projects a tree x onto the vector defined by a tree a:

\[\operatorname{proj}_a x = \frac{\langle x, a \rangle}{\langle a, a \rangle} a\]

Parameters:

• x – tree to project.

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

Returns:

tree with the same structure as x.