Wrong implementation of hyper_v2 mix_layers

Hi,
In VeLO (https://arxiv.org/pdf/2211.09760.pdf) Section B.3, it states that mixing is done by F0(x) + max(σ(F1(σ(F2(x)))), axis = 0, keep_dims = True).
However, in the implementation of hyper_v2 (https://github.com/google/learned_optimization/blob/main/learned_optimization/research/general_lopt/hyper_v2.py#L330-L335), it essentially use only one linear layer instead of two, as the input to second linear layer is x instead of mix_layer (L332).

Start with Section B.3 of the linked VeLO paper and inspect learned_optimization/research/general_lopt/hyper_v2.py around lines 330-335. Trace the inputs to the linear layers and compare them with the paper's mix_layers expression. Done means the implementation's layer inputs match the documented formula and relevant existing tests pass.

Scritto dal modello di indicizzazione a partire dal testo della issue.