Fixes the aux output being from the wrong y.

optimistix/_solver/gauss_newton.py

@@ -17,6 +17,7 @@
     max_norm,
     sum_squares,
     tree_full_like,
+    tree_where,
     verbose_print,
 )
 from .._search import (
@@ -331,6 +332,7 @@ def rejected(descent_state):
             search_result == RESULTS.successful, descent_result, search_result
         )
 
+        prev_aux = tree_where(state.first_step, aux, state.aux)
         state = _GaussNewtonState(
             first_step=jnp.array(False),
             y_eval=y_eval,
@@ -344,7 +346,7 @@ def rejected(descent_state):
             num_accepted_steps=num_accepted_steps,
             num_steps_since_acceptance=num_steps_since_acceptance,
         )
-        return y, state, aux
+        return y, state, prev_aux
 
     def terminate(
         self,

optimistix/_solver/gradient_methods.py

@@ -16,6 +16,7 @@
     lin_to_grad,
     max_norm,
     tree_full_like,
+    tree_where,
 )
 from .._search import (
     AbstractDescent,
@@ -207,6 +208,7 @@ def rejected(descent_state):
             search_result == RESULTS.successful, descent_result, search_result
         )
 
+        prev_aux = tree_where(state.first_step, aux, state.aux)
         state = _GradientDescentState(
             first_step=jnp.array(False),
             y_eval=y_eval,
@@ -217,7 +219,7 @@ def rejected(descent_state):
             terminate=terminate,
             result=result,
         )
-        return y, state, aux
+        return y, state, prev_aux
 
     def terminate(
         self,

optimistix/_solver/quasi_newton.py

@@ -20,6 +20,7 @@
     max_norm,
     tree_dot,
     tree_full_like,
+    tree_where,
     verbose_print,
 )
 from .._search import (
@@ -286,6 +287,7 @@ def rejected(descent_state):
             search_result == RESULTS.successful, descent_result, search_result
         )
 
+        prev_aux = tree_where(state.first_step, aux, state.aux)
         state = _QuasiNewtonState(
             first_step=jnp.array(False),
             y_eval=y_eval,
@@ -298,7 +300,7 @@ def rejected(descent_state):
             num_accepted_steps=state.num_accepted_steps + jnp.where(accept, 1, 0),
             hessian_update_state=hessian_update_state,
         )
-        return y, state, aux
+        return y, state, prev_aux
 
     def terminate(
         self,

tests/test_best_so_far.py

@@ -1,5 +1,6 @@
 import jax.numpy as jnp
 import optimistix as optx
+import pytest
 
 
 def test_fixed_point():
@@ -41,3 +42,25 @@ def fn(y, _):
     solver = optx.BestSoFarMinimiser(solver)
     sol = optx.minimise(fn, solver, jnp.array(0.0))
     assert jnp.allclose(sol.value, 0.96118069, rtol=1e-5, atol=1e-5)
+
+
+# https://github.com/patrick-kidger/optimistix/issues/33
+@pytest.mark.parametrize(
+    "solver", (optx.BFGS(rtol=1e-5, atol=1e-5), optx.NonlinearCG(atol=1e-5, rtol=1e-5))
+)
+def test_checks_last_point_minimiser(solver):
+    def fn(y, _):
+        return (y - 3.0) ** 2
+
+    solver = optx.BestSoFarMinimiser(solver)
+    sol = optx.minimise(fn, solver, jnp.array(0.0))
+    assert sol.value == 3.0
+
+
+def test_checks_last_point_least_squares():
+    def fn(y, _):
+        return y - 3.0
+
+    solver = optx.BestSoFarLeastSquares(optx.GaussNewton(rtol=1e-5, atol=1e-5))
+    sol = optx.least_squares(fn, solver, jnp.array(0.0))
+    assert sol.value == 3.0