From 857359a77a45816873110fb2b6d5ca4a372a300e Mon Sep 17 00:00:00 2001 From: Bence Bagi Date: Mon, 18 May 2026 14:45:48 +0200 Subject: [PATCH] Optionally pass more info at y_eval to searches --- optimistix/_search.py | 38 ++++++-------------- optimistix/_solver/_f_eval_info_helpers.py | 42 ++++++++++++++++++++++ optimistix/_solver/backtracking.py | 3 +- optimistix/_solver/gradient_methods.py | 34 +++++++++++++----- optimistix/_solver/learning_rate.py | 3 +- optimistix/_solver/nonlinear_cg.py | 9 +++-- optimistix/_solver/quasi_newton.py | 28 ++++++++++----- optimistix/_solver/trust_region.py | 4 ++- 8 files changed, 112 insertions(+), 49 deletions(-) create mode 100644 optimistix/_solver/_f_eval_info_helpers.py diff --git a/optimistix/_search.py b/optimistix/_search.py index 2f4ec172..09407923 100644 --- a/optimistix/_search.py +++ b/optimistix/_search.py @@ -29,6 +29,7 @@ import jax.numpy as jnp import jax.tree_util as jtu import lineax as lx +from equinox import AbstractClassVar from jaxtyping import Array, Bool, Scalar from ._custom_types import ( @@ -79,56 +80,39 @@ def as_min(self): # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right. -class EvalGrad(FunctionInfo, Generic[Y]): +class EvalGrad(Eval, Generic[Y]): """Has a `.f` attribute as with [`optimistix.FunctionInfo.Eval`][]. Also has a `.grad` attribute describing `d(fn)/dy`. Used with first-order solvers for minimisation problems. (E.g. gradient descent; nonlinear CG.) """ - f: Scalar grad: Y - def as_min(self): - return self.f + def compute_grad(self): + return self.grad def compute_grad_dot(self, y: Y): return tree_dot(self.grad, y) # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right. -class EvalGradHessian(FunctionInfo, Generic[Y]): +class EvalGradHessian(EvalGrad[Y], Generic[Y]): """Has `.f` and `.grad` attributes as with [`optimistix.FunctionInfo.EvalGrad`][]. Also has a `.hessian` attribute describing (an approximation to) the Hessian of `fn` at `y`. Used with quasi-Newton minimisation algorithms, like BFGS. """ - f: Scalar - grad: Y hessian: lx.AbstractLinearOperator - def as_min(self): - return self.f - - def compute_grad_dot(self, y: Y): - return tree_dot(self.grad, y) - # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right. -class EvalGradHessianInv(FunctionInfo, Generic[Y]): +class EvalGradHessianInv(EvalGrad[Y], Generic[Y]): """As [`optimistix.FunctionInfo.EvalGradHessian`][], but records the (approximate) inverse-Hessian instead. Has `.f` and `.grad` and `.hessian_inv` attributes. """ - f: Scalar - grad: Y hessian_inv: lx.AbstractLinearOperator - def as_min(self): - return self.f - - def compute_grad_dot(self, y: Y): - return tree_dot(self.grad, y) - # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right. class Residual(FunctionInfo, Generic[Out]): @@ -143,18 +127,14 @@ def as_min(self): # NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right. -class ResidualJac(FunctionInfo, Generic[Y, Out]): +class ResidualJac(Residual[Out], Generic[Y, Out]): """Records the Jacobian `d(fn)/dy` as a linear operator. Used for least squares problems, for which `fn` returns residuals. Has `.residual` and `.jac` attributes, where `residual = fn(y)`, `jac = d(fn)/dy`. """ - residual: Out jac: lx.AbstractLinearOperator - def as_min(self): - return 0.5 * sum_squares(self.residual) - def compute_grad(self): conj_residual = jtu.tree_map(jnp.conj, self.residual) return self.jac.transpose().mv(conj_residual) @@ -320,6 +300,8 @@ class AbstractSearch(eqx.Module, Generic[Y, _FnInfo, _FnEvalInfo, SearchState]): See [this documentation](./introduction.md) for more information. """ + info_needed_at_y_eval: AbstractClassVar[type[FunctionInfo]] + @abc.abstractmethod def init(self, y: Y, f_info_struct: _FnInfo) -> SearchState: """Is called just once, at the very start of the entire optimisation problem. @@ -363,7 +345,7 @@ def step( - `f_info`: An [`optimistix.FunctionInfo`][] describing information about `f` evaluated at `y`, the gradient of `f` at `y`, etc. - `f_eval_info`: An [`optimistix.FunctionInfo`][] describing information about - `f` evaluated at `y`, the gradient of `f` at `y`, etc. + `f` evaluated at `y_eval`, the gradient of `f` at `y_eval`, etc. - `state`: the evolving state of the repeated searches. **Returns:** diff --git a/optimistix/_solver/_f_eval_info_helpers.py b/optimistix/_solver/_f_eval_info_helpers.py new file mode 100644 index 00000000..b62aaa33 --- /dev/null +++ b/optimistix/_solver/_f_eval_info_helpers.py @@ -0,0 +1,42 @@ +from collections.abc import Callable + +from jaxtyping import Scalar + +from .._custom_types import Y +from .._misc import lin_to_grad +from .._search import FunctionInfo + + +def make_eval_info_for_search( + info_needed: type[FunctionInfo], + y_eval: Y, + f_eval: Scalar, + lin_fn: Callable[[Y], Scalar], + autodiff_mode: str, +) -> FunctionInfo.Eval | FunctionInfo.EvalGrad: + """Build the cheapest scalar `FunctionInfo` satisfying `info_needed`.""" + if issubclass(FunctionInfo.Eval, info_needed): + return FunctionInfo.Eval(f_eval) + elif issubclass(FunctionInfo.EvalGrad, info_needed): + grad = lin_to_grad(lin_fn, y_eval, autodiff_mode, f_eval.dtype) + return FunctionInfo.EvalGrad(f_eval, grad) + else: + raise ValueError( + f"Cannot provide requested function information at `y_eval`: {info_needed}." + " Only FunctionInfo.Eval and FunctionInfo.EvalGrad are supported." + ) + + +def promote_to_eval_grad( + f_eval_info: FunctionInfo.Eval | FunctionInfo.EvalGrad, + y_eval: Y, + f_eval: Scalar, + lin_fn: Callable[[Y], Scalar], + autodiff_mode: str, +) -> FunctionInfo.EvalGrad: + """Reuse an existing scalar gradient evaluation, or compute one if needed.""" + if isinstance(f_eval_info, FunctionInfo.EvalGrad): + return f_eval_info + else: + grad = lin_to_grad(lin_fn, y_eval, autodiff_mode, f_eval.dtype) + return FunctionInfo.EvalGrad(f_eval, grad) diff --git a/optimistix/_solver/backtracking.py b/optimistix/_solver/backtracking.py index a9cdbf06..1bb53c82 100644 --- a/optimistix/_solver/backtracking.py +++ b/optimistix/_solver/backtracking.py @@ -1,4 +1,4 @@ -from typing import cast, TypeAlias +from typing import cast, ClassVar, TypeAlias import equinox as eqx import jax.numpy as jnp @@ -29,6 +29,7 @@ class BacktrackingArmijo(AbstractSearch[Y, _FnInfo, _FnEvalInfo, _BacktrackingSt decrease_factor: ScalarLike = 0.5 slope: ScalarLike = 0.1 step_init: ScalarLike = 1.0 + info_needed_at_y_eval: ClassVar[type[FunctionInfo]] = FunctionInfo def __post_init__(self): self.decrease_factor = eqx.error_if( diff --git a/optimistix/_solver/gradient_methods.py b/optimistix/_solver/gradient_methods.py index 67a86498..b7ac5bde 100644 --- a/optimistix/_solver/gradient_methods.py +++ b/optimistix/_solver/gradient_methods.py @@ -13,7 +13,6 @@ from .._misc import ( cauchy_termination, filter_cond, - lin_to_grad, max_norm, tree_full_like, tree_where, @@ -24,6 +23,7 @@ FunctionInfo, ) from .._solution import RESULTS +from ._f_eval_info_helpers import make_eval_info_for_search, promote_to_eval_grad from .learning_rate import LearningRate @@ -132,7 +132,9 @@ class AbstractGradientDescent(AbstractMinimiser[Y, Aux, _GradientDescentState]): norm: AbstractVar[Callable[[PyTree], Scalar]] descent: AbstractVar[AbstractDescent[Y, FunctionInfo.EvalGrad, Any]] search: AbstractVar[ - AbstractSearch[Y, FunctionInfo.EvalGrad, FunctionInfo.Eval, Any] + AbstractSearch[ + Y, FunctionInfo.EvalGrad, FunctionInfo.Eval | FunctionInfo.EvalGrad, Any + ] ] def init( @@ -171,20 +173,30 @@ def step( f_eval, lin_fn, aux_eval = jax.linearize( lambda _y: fn(_y, args), state.y_eval, has_aux=True ) + f_eval_info = make_eval_info_for_search( + self.search.info_needed_at_y_eval, + state.y_eval, + f_eval, + lin_fn, + autodiff_mode, + ) + step_size, accept, search_result, search_state = self.search.step( state.first_step, y, state.y_eval, state.f_info, - FunctionInfo.Eval(f_eval), + f_eval_info, state.search_state, ) def accepted(descent_state): - grad = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype) - - f_eval_info = FunctionInfo.EvalGrad(f_eval, grad) - descent_state = self.descent.query(state.y_eval, f_eval_info, descent_state) + accepted_f_eval_info = promote_to_eval_grad( + f_eval_info, state.y_eval, f_eval, lin_fn, autodiff_mode + ) + descent_state = self.descent.query( + state.y_eval, accepted_f_eval_info, descent_state + ) y_diff = (state.y_eval**ω - y**ω).ω f_diff = (f_eval**ω - state.f_info.f**ω).ω terminate = cauchy_termination( @@ -193,7 +205,13 @@ def accepted(descent_state): terminate = jnp.where( state.first_step, jnp.array(False), terminate ) # Skip termination on first step - return state.y_eval, f_eval_info, aux_eval, descent_state, terminate + return ( + state.y_eval, + accepted_f_eval_info, + aux_eval, + descent_state, + terminate, + ) def rejected(descent_state): return y, state.f_info, state.aux, descent_state, jnp.array(False) diff --git a/optimistix/_solver/learning_rate.py b/optimistix/_solver/learning_rate.py index 36a0a22e..55d6ea0c 100644 --- a/optimistix/_solver/learning_rate.py +++ b/optimistix/_solver/learning_rate.py @@ -1,4 +1,4 @@ -from typing import cast +from typing import cast, ClassVar import equinox as eqx import jax.numpy as jnp @@ -16,6 +16,7 @@ def _typed_asarray(x: ScalarLike) -> Array: class LearningRate(AbstractSearch[Y, FunctionInfo, FunctionInfo, None]): """Move downhill by taking a step of the fixed size `learning_rate`.""" + info_needed_at_y_eval: ClassVar[type[FunctionInfo]] = FunctionInfo learning_rate: ScalarLike = eqx.field(converter=_typed_asarray) def init(self, y: Y, f_info_struct: FunctionInfo) -> None: diff --git a/optimistix/_solver/nonlinear_cg.py b/optimistix/_solver/nonlinear_cg.py index 73ed02bd..88ff8d0b 100644 --- a/optimistix/_solver/nonlinear_cg.py +++ b/optimistix/_solver/nonlinear_cg.py @@ -183,7 +183,9 @@ class NonlinearCG(AbstractGradientDescent[Y, Aux]): atol: float norm: Callable[[PyTree], Scalar] descent: NonlinearCGDescent[Y] - search: AbstractSearch[Y, FunctionInfo.EvalGrad, FunctionInfo.Eval, Any] + search: AbstractSearch[ + Y, FunctionInfo.EvalGrad, FunctionInfo.Eval | FunctionInfo.EvalGrad, Any + ] def __init__( self, @@ -193,7 +195,10 @@ def __init__( method: Callable[[Y, Y, Y], Scalar] = polak_ribiere, # TODO(raderj): replace the default line search with something better. search: AbstractSearch[ - Y, FunctionInfo.EvalGrad, FunctionInfo.Eval, Any + Y, + FunctionInfo.EvalGrad, + FunctionInfo.Eval | FunctionInfo.EvalGrad, + Any, ] = BacktrackingArmijo(decrease_factor=0.5, slope=0.1), ): """**Arguments:** diff --git a/optimistix/_solver/quasi_newton.py b/optimistix/_solver/quasi_newton.py index 94f38810..16541446 100644 --- a/optimistix/_solver/quasi_newton.py +++ b/optimistix/_solver/quasi_newton.py @@ -17,7 +17,6 @@ cauchy_termination, default_verbose, filter_cond, - lin_to_grad, max_norm, tree_dot, tree_full_like, @@ -29,6 +28,7 @@ FunctionInfo, ) from .._solution import RESULTS +from ._f_eval_info_helpers import make_eval_info_for_search, promote_to_eval_grad from .backtracking import BacktrackingArmijo from .gauss_newton import NewtonDescent @@ -139,7 +139,9 @@ class AbstractQuasiNewton( norm: AbstractVar[Callable[[PyTree], Scalar]] use_inverse: AbstractVar[bool] descent: AbstractVar[AbstractDescent[Y, _Hessian, Any]] - search: AbstractVar[AbstractSearch[Y, _Hessian, FunctionInfo.Eval, Any]] + search: AbstractVar[ + AbstractSearch[Y, _Hessian, FunctionInfo.Eval | FunctionInfo.EvalGrad, Any] + ] verbose: AbstractVar[Callable[..., None]] @abc.abstractmethod @@ -214,29 +216,39 @@ def step( f_eval, lin_fn, aux_eval = jax.linearize( lambda _y: fn(_y, args), state.y_eval, has_aux=True ) + f_eval_info = make_eval_info_for_search( + self.search.info_needed_at_y_eval, + state.y_eval, + f_eval, + lin_fn, + autodiff_mode, + ) + step_size, accept, search_result, search_state = self.search.step( state.first_step, y, state.y_eval, state.f_info, - FunctionInfo.Eval(f_eval), + f_eval_info, state.search_state, ) def accepted(descent_state): - grad = lin_to_grad(lin_fn, state.y_eval, autodiff_mode, f_eval.dtype) + accepted_f_eval_info = promote_to_eval_grad( + f_eval_info, state.y_eval, f_eval, lin_fn, autodiff_mode + ) - f_eval_info, hessian_update_state = self.update_hessian( + updated_f_eval_info, hessian_update_state = self.update_hessian( y, state.y_eval, state.f_info, - FunctionInfo.EvalGrad(f_eval, grad), + accepted_f_eval_info, state.hessian_update_state, ) descent_state = self.descent.query( state.y_eval, - f_eval_info, # pyright: ignore + updated_f_eval_info, descent_state, ) y_diff = (state.y_eval**ω - y**ω).ω @@ -249,7 +261,7 @@ def accepted(descent_state): ) # Skip termination on first step return ( state.y_eval, - f_eval_info, + updated_f_eval_info, aux_eval, descent_state, terminate, diff --git a/optimistix/_solver/trust_region.py b/optimistix/_solver/trust_region.py index 15d171d2..0c039413 100644 --- a/optimistix/_solver/trust_region.py +++ b/optimistix/_solver/trust_region.py @@ -1,5 +1,5 @@ import abc -from typing import TypeAlias, TypeVar +from typing import ClassVar, TypeAlias, TypeVar import equinox as eqx import jax.numpy as jnp @@ -47,6 +47,8 @@ class _AbstractTrustRegion(AbstractSearch[Y, _FnInfo, _FnEvalInfo, _TrustRegionS high_constant: AbstractVar[ScalarLike] low_constant: AbstractVar[ScalarLike] + info_needed_at_y_eval: ClassVar[type[FunctionInfo]] = FunctionInfo + def __post_init__(self): # You would not expect `self.low_cutoff` or `self.high_cutoff` to # be below zero, but this is technically not incorrect so we don't