Fix the jax line search, the zoom may only store two points in lo, hi and rec.

PiperOrigin-RevId: 833328457

Author: Google-ML-Automation

jax/_src/scipy/optimize/line_search.py

@@ -191,6 +191,16 @@ def body(state):
             ),
         ),
     )
+    state = state._replace(
+        **_binary_replace(
+            lo_to_j & ~hi_to_lo,
+            state._asdict(),
+            dict(
+                a_rec=state.a_lo,
+                phi_rec=state.phi_lo,
+            ),
+        ),
+    )
     state = state._replace(
         **_binary_replace(
             lo_to_j,
@@ -199,8 +209,6 @@ def body(state):
                 a_lo=a_j,
                 phi_lo=phi_j,
                 dphi_lo=dphi_j,
-                a_rec=state.a_lo,
-                phi_rec=state.phi_lo,
             ),
         ),
     )