torchzero 0.3.11__py3-none-any.whl → 0.3.14__py3-none-any.whl

torchzero/modules/line_search/backtracking.py

@@ -4,7 +4,7 @@ from operator import itemgetter
 
 import torch
 
-from .line_search import LineSearchBase
+from .line_search import LineSearchBase, TerminationCondition, termination_condition
 
 
 def backtracking_line_search(
@@ -14,7 +14,7 @@ def backtracking_line_search(
     beta: float = 0.5,
     c: float = 1e-4,
     maxiter: int = 10,
-    try_negative: bool = False,
+    condition: TerminationCondition = 'armijo',
 ) -> float | None:
     """
 
@@ -31,16 +31,20 @@ def backtracking_line_search(
     """
 
     a = init
-    f_x = f(0)
+    f_0 = f(0)
     f_prev = None
 
     for iteration in range(maxiter):
         f_a = f(a)
+        if not math.isfinite(f_a):
+            a *= beta
+            continue
 
-        if (f_prev is not None) and (f_a > f_prev) and (f_prev < f_x): return a / beta
+        if (f_prev is not None) and (f_a > f_prev) and (f_prev < f_0):
+            return a / beta # new value is larger than previous value
         f_prev = f_a
 
-        if f_a < f_x + c * a * min(g_0, 0): # pyright: ignore[reportArgumentType]
+        if termination_condition(condition, f_0=f_0, g_0=g_0, f_a=f_a, g_a=None, a=a, c=c):
             # found an acceptable alpha
             return a
 
@@ -48,53 +52,45 @@ def backtracking_line_search(
         a *= beta
 
     # fail
-    if try_negative:
-        def inv_objective(alpha): return f(-alpha)
-
-        v = backtracking_line_search(
-            inv_objective,
-            g_0=-g_0,
-            beta=beta,
-            c=c,
-            maxiter=maxiter,
-            try_negative=False,
-        )
-        if v is not None: return -v
-
     return None
 
 class Backtracking(LineSearchBase):
-    """Backtracking line search satisfying the Armijo condition.
+    """Backtracking line search.
 
     Args:
         init (float, optional): initial step size. Defaults to 1.0.
         beta (float, optional): multiplies each consecutive step size by this value. Defaults to 0.5.
-        c (float, optional): acceptance value for Armijo condition. Defaults to 1e-4.
-        maxiter (int, optional): Maximum line search function evaluations. Defaults to 10.
+        c (float, optional): sufficient decrease condition. Defaults to 1e-4.
+        condition (TerminationCondition, optional):
+            termination condition, only ones that do not use gradient at f(x+a*d) can be specified.
+            - "armijo" - sufficient decrease condition.
+            - "decrease" - any decrease in objective function value satisfies the condition.
+
+            "goldstein" can techincally be specified but it doesn't make sense because there is not zoom stage.
+            Defaults to 'armijo'.
+        maxiter (int, optional): maximum number of function evaluations per step. Defaults to 10.
         adaptive (bool, optional):
-            when enabled, if line search failed, beta is reduced.
-            Otherwise it is reset to initial value. Defaults to True.
-        try_negative (bool, optional): Whether to perform line search in opposite direction on fail. Defaults to False.
+            when enabled, if line search failed, step size will continue decreasing on the next step.
+            Otherwise it will restart the line search from ``init`` step size. Defaults to True.
 
     Examples:
-        Gradient descent with backtracking line search:
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Backtracking()
-            )
-
-        LBFGS with backtracking line search:
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.LBFGS(),
-                tz.m.Backtracking()
-            )
+    Gradient descent with backtracking line search:
+
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Backtracking()
+    )
+    ```
+
+    L-BFGS with backtracking line search:
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.LBFGS(),
+        tz.m.Backtracking()
+    )
+    ```
 
     """
     def __init__(
@@ -102,41 +98,47 @@ class Backtracking(LineSearchBase):
         init: float = 1.0,
         beta: float = 0.5,
         c: float = 1e-4,
+        condition: TerminationCondition = 'armijo',
         maxiter: int = 10,
         adaptive=True,
-        try_negative: bool = False,
     ):
-        defaults=dict(init=init,beta=beta,c=c,maxiter=maxiter,adaptive=adaptive, try_negative=try_negative)
+        defaults=dict(init=init,beta=beta,c=c,condition=condition,maxiter=maxiter,adaptive=adaptive)
         super().__init__(defaults=defaults)
-        self.global_state['beta_scale'] = 1.0
 
     def reset(self):
         super().reset()
-        self.global_state['beta_scale'] = 1.0
 
     @torch.no_grad
     def search(self, update, var):
-        init, beta, c, maxiter, adaptive, try_negative = itemgetter(
-            'init', 'beta', 'c', 'maxiter', 'adaptive', 'try_negative')(self.settings[var.params[0]])
+        init, beta, c, condition, maxiter, adaptive = itemgetter(
+            'init', 'beta', 'c', 'condition', 'maxiter', 'adaptive')(self.defaults)
 
         objective = self.make_objective(var=var)
 
         # # directional derivative
-        d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), var.get_update()))
+        if c == 0: d = 0
+        else: d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), var.get_update()))
 
-        # scale beta (beta is multiplicative and i think may be better than scaling initial step size)
-        if adaptive: beta = beta * self.global_state['beta_scale']
+        # scale init
+        init_scale = self.global_state.get('init_scale', 1)
+        if adaptive: init = init * init_scale
 
-        step_size = backtracking_line_search(objective, d, init=init,beta=beta,
-                                        c=c,maxiter=maxiter, try_negative=try_negative)
+        step_size = backtracking_line_search(objective, d, init=init, beta=beta,c=c, condition=condition, maxiter=maxiter)
 
         # found an alpha that reduces loss
         if step_size is not None:
-            self.global_state['beta_scale'] = min(1.0, self.global_state['beta_scale'] * math.sqrt(1.5))
+            #self.global_state['beta_scale'] = min(1.0, self.global_state['beta_scale'] * math.sqrt(1.5))
+            self.global_state['init_scale'] = 1
             return step_size
 
-        # on fail reduce beta scale value
-        self.global_state['beta_scale'] /= 1.5
+        # on fail set init_scale to continue decreasing the step size
+        # or set to large step size when it becomes too small
+        if adaptive:
+            finfo = torch.finfo(var.params[0].dtype)
+            if init_scale <= finfo.tiny * 2:
+                self.global_state["init_scale"] = finfo.max / 2
+            else:
+                self.global_state['init_scale'] = init_scale * beta**maxiter
         return 0
 
 def _lerp(start,end,weight):
@@ -147,30 +149,37 @@ class AdaptiveBacktracking(LineSearchBase):
     such that optimal step size in the procedure would be found on the second line search iteration.
 
     Args:
-        init (float, optional): step size for the first step. Defaults to 1.0.
+        init (float, optional): initial step size. Defaults to 1.0.
         beta (float, optional): multiplies each consecutive step size by this value. Defaults to 0.5.
-        c (float, optional): acceptance value for Armijo condition. Defaults to 1e-4.
-        maxiter (int, optional): Maximum line search function evaluations. Defaults to 10.
+        c (float, optional): sufficient decrease condition. Defaults to 1e-4.
+        condition (TerminationCondition, optional):
+            termination condition, only ones that do not use gradient at f(x+a*d) can be specified.
+            - "armijo" - sufficient decrease condition.
+            - "decrease" - any decrease in objective function value satisfies the condition.
+
+            "goldstein" can techincally be specified but it doesn't make sense because there is not zoom stage.
+            Defaults to 'armijo'.
+        maxiter (int, optional): maximum number of function evaluations per step. Defaults to 10.
         target_iters (int, optional):
-            target number of iterations that would be performed until optimal step size is found. Defaults to 1.
+            sets next step size such that this number of iterations are expected
+            to be performed until optimal step size is found. Defaults to 1.
         nplus (float, optional):
-            Multiplier to initial step size if it was found to be the optimal step size. Defaults to 2.0.
+            if initial step size is optimal, it is multiplied by this value. Defaults to 2.0.
         scale_beta (float, optional):
-            Momentum for initial step size, at 0 disables momentum. Defaults to 0.0.
-        try_negative (bool, optional): Whether to perform line search in opposite direction on fail. Defaults to False.
+            momentum for initial step size, at 0 disables momentum. Defaults to 0.0.
     """
     def __init__(
         self,
         init: float = 1.0,
         beta: float = 0.5,
         c: float = 1e-4,
+        condition: TerminationCondition = 'armijo',
         maxiter: int = 20,
         target_iters = 1,
         nplus = 2.0,
         scale_beta = 0.0,
-        try_negative: bool = False,
     ):
-        defaults=dict(init=init,beta=beta,c=c,maxiter=maxiter,target_iters=target_iters,nplus=nplus,scale_beta=scale_beta, try_negative=try_negative)
+        defaults=dict(init=init,beta=beta,c=c,condition=condition,maxiter=maxiter,target_iters=target_iters,nplus=nplus,scale_beta=scale_beta)
         super().__init__(defaults=defaults)
 
         self.global_state['beta_scale'] = 1.0
@@ -183,8 +192,8 @@ class AdaptiveBacktracking(LineSearchBase):
 
     @torch.no_grad
     def search(self, update, var):
-        init, beta, c, maxiter, target_iters, nplus, scale_beta, try_negative=itemgetter(
-            'init','beta','c','maxiter','target_iters','nplus','scale_beta', 'try_negative')(self.settings[var.params[0]])
+        init, beta, c,condition, maxiter, target_iters, nplus, scale_beta=itemgetter(
+            'init','beta','c','condition', 'maxiter','target_iters','nplus','scale_beta')(self.defaults)
 
         objective = self.make_objective(var=var)
 
@@ -198,8 +207,7 @@ class AdaptiveBacktracking(LineSearchBase):
         # scale step size so that decrease is expected at target_iters
         init = init * self.global_state['initial_scale']
 
-        step_size = backtracking_line_search(objective, d, init=init, beta=beta,
-                                        c=c,maxiter=maxiter, try_negative=try_negative)
+        step_size = backtracking_line_search(objective, d, init=init, beta=beta, c=c, condition=condition, maxiter=maxiter)
 
         # found an alpha that reduces loss
         if step_size is not None:
@@ -208,7 +216,12 @@ class AdaptiveBacktracking(LineSearchBase):
             # initial step size satisfied conditions, increase initial_scale by nplus
             if step_size == init and target_iters > 0:
                 self.global_state['initial_scale'] *= nplus ** target_iters
-                self.global_state['initial_scale'] = min(self.global_state['initial_scale'], 1e32) # avoid overflow error
+
+                # clip by maximum possibel value to avoid overflow exception
+                self.global_state['initial_scale'] = min(
+                    self.global_state['initial_scale'],
+                    torch.finfo(var.params[0].dtype).max / 2,
+                )
 
             else:
                 # otherwise make initial_scale such that target_iters iterations will satisfy armijo

torchzero/modules/line_search/line_search.py

@@ -3,13 +3,13 @@ from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from functools import partial
 from operator import itemgetter
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
 import torch
 
 from ...core import Module, Target, Var
-from ...utils import tofloat
+from ...utils import tofloat, set_storage_
 
 
 class MaxLineSearchItersReached(Exception): pass
@@ -29,60 +29,59 @@ class LineSearchBase(Module, ABC):
             doesn't have a maxiter option. Defaults to None.
 
     Other useful methods:
-        * `evaluate_step_size` - returns loss with a given scalar step size
-        * `evaluate_step_size_loss_and_derivative` - returns loss and directional derivative with a given scalar step size
-        * `make_objective` - creates a function that accepts a scalar step size and returns loss. This can be passed to a scalar solver, such as scipy.optimize.minimize_scalar.
-        * `make_objective_with_derivative` - creates a function that accepts a scalar step size and returns a tuple with loss and directional derivative. This can be passed to a scalar solver.
+        * ``evaluate_f`` - returns loss with a given scalar step size
+        * ``evaluate_f_d`` - returns loss and directional derivative with a given scalar step size
+        * ``make_objective`` - creates a function that accepts a scalar step size and returns loss. This can be passed to a scalar solver, such as scipy.optimize.minimize_scalar.
+        * ``make_objective_with_derivative`` - creates a function that accepts a scalar step size and returns a tuple with loss and directional derivative. This can be passed to a scalar solver.
 
     Examples:
-        #### Basic line search
 
-        This evaluates all step sizes in a range by using the :code:`self.evaluate_step_size` method.
+    #### Basic line search
 
-        .. code-block:: python
+    This evaluates all step sizes in a range by using the :code:`self.evaluate_step_size` method.
+    ```python
+    class GridLineSearch(LineSearch):
+        def __init__(self, start, end, num):
+            defaults = dict(start=start,end=end,num=num)
+            super().__init__(defaults)
 
-            class GridLineSearch(LineSearch):
-                def __init__(self, start, end, num):
-                    defaults = dict(start=start,end=end,num=num)
-                    super().__init__(defaults)
+        @torch.no_grad
+        def search(self, update, var):
 
-                @torch.no_grad
-                def search(self, update, var):
-                    settings = self.settings[var.params[0]]
-                    start = settings["start"]
-                    end = settings["end"]
-                    num = settings["num"]
+            start = self.defaults["start"]
+            end = self.defaults["end"]
+            num = self.defaults["num"]
 
-                    lowest_loss = float("inf")
-                    best_step_size = best_step_size
+            lowest_loss = float("inf")
+            best_step_size = best_step_size
 
-                    for step_size in torch.linspace(start,end,num):
-                        loss = self.evaluate_step_size(step_size.item(), var=var, backward=False)
-                        if loss < lowest_loss:
-                            lowest_loss = loss
-                            best_step_size = step_size
+            for step_size in torch.linspace(start,end,num):
+                loss = self.evaluate_step_size(step_size.item(), var=var, backward=False)
+                if loss < lowest_loss:
+                    lowest_loss = loss
+                    best_step_size = step_size
 
-                    return best_step_size
+            return best_step_size
+    ```
 
-        #### Using external solver via self.make_objective
+    #### Using external solver via self.make_objective
 
-        Here we let :code:`scipy.optimize.minimize_scalar` solver find the best step size via :code:`self.make_objective`
+    Here we let :code:`scipy.optimize.minimize_scalar` solver find the best step size via :code:`self.make_objective`
 
-        .. code-block:: python
+    ```python
+    class ScipyMinimizeScalar(LineSearch):
+        def __init__(self, method: str | None = None):
+            defaults = dict(method=method)
+            super().__init__(defaults)
 
-            class ScipyMinimizeScalar(LineSearch):
-                def __init__(self, method: str | None = None):
-                    defaults = dict(method=method)
-                    super().__init__(defaults)
-
-                @torch.no_grad
-                def search(self, update, var):
-                    objective = self.make_objective(var=var)
-                    method = self.settings[var.params[0]]["method"]
-
-                    res = self.scopt.minimize_scalar(objective, method=method)
-                    return res.x
+        @torch.no_grad
+        def search(self, update, var):
+            objective = self.make_objective(var=var)
+            method = self.defaults["method"]
 
+            res = self.scopt.minimize_scalar(objective, method=method)
+            return res.x
+    ```
     """
     def __init__(self, defaults: dict[str, Any] | None, maxiter: int | None = None):
         super().__init__(defaults)
@@ -94,6 +93,7 @@ class LineSearchBase(Module, ABC):
         self._lowest_loss = float('inf')
         self._best_step_size: float = 0
         self._current_iter = 0
+        self._initial_params = None
 
     def set_step_size_(
         self,
@@ -102,10 +102,27 @@ class LineSearchBase(Module, ABC):
         update: list[torch.Tensor],
     ):
         if not math.isfinite(step_size): return
-        step_size = max(min(tofloat(step_size), 1e36), -1e36) # fixes overflow when backtracking keeps increasing alpha after converging
-        alpha = self._current_step_size - step_size
-        if alpha != 0:
-            torch._foreach_add_(params, update, alpha=alpha)
+
+         # fixes overflow when backtracking keeps increasing alpha after converging
+        step_size = max(min(tofloat(step_size), 1e36), -1e36)
+
+        # skip is parameters are already at suggested step size
+        if self._current_step_size == step_size: return
+
+        # this was basically causing floating point imprecision to build up
+        #if False:
+        # if abs(alpha) < abs(step_size) and step_size != 0:
+        #     torch._foreach_add_(params, update, alpha=alpha)
+
+        # else:
+        assert self._initial_params is not None
+        if step_size == 0:
+            new_params = [p.clone() for p in self._initial_params]
+        else:
+            new_params = torch._foreach_sub(self._initial_params, update, alpha=step_size)
+        for c, n in zip(params, new_params):
+            set_storage_(c, n)
+
         self._current_step_size = step_size
 
     def _set_per_parameter_step_size_(
@@ -114,10 +131,20 @@ class LineSearchBase(Module, ABC):
         params: list[torch.Tensor],
         update: list[torch.Tensor],
     ):
-        if not np.isfinite(step_size): step_size = [0 for _ in step_size]
-        alpha = [self._current_step_size - s for s in step_size]
-        if any(a!=0 for a in alpha):
-            torch._foreach_add_(params, torch._foreach_mul(update, alpha))
+        # if not np.isfinite(step_size): step_size = [0 for _ in step_size]
+        # alpha = [self._current_step_size - s for s in step_size]
+        # if any(a!=0 for a in alpha):
+        #     torch._foreach_add_(params, torch._foreach_mul(update, alpha))
+        assert self._initial_params is not None
+        if not np.isfinite(step_size).all(): step_size = [0 for _ in step_size]
+
+        if any(s!=0 for s in step_size):
+            new_params = torch._foreach_sub(self._initial_params, torch._foreach_mul(update, step_size))
+        else:
+            new_params = [p.clone() for p in self._initial_params]
+
+        for c, n in zip(params, new_params):
+            set_storage_(c, n)
 
     def _loss(self, step_size: float, var: Var, closure, params: list[torch.Tensor],
               update: list[torch.Tensor], backward:bool=False) -> float:
@@ -149,7 +176,7 @@ class LineSearchBase(Module, ABC):
 
         return tofloat(loss)
 
-    def _loss_derivative(self, step_size: float, var: Var, closure,
+    def _loss_derivative_gradient(self, step_size: float, var: Var, closure,
                          params: list[torch.Tensor], update: list[torch.Tensor]):
         # if step_size is 0, we might already know the derivative
         if (var.grad is not None) and (step_size == 0):
@@ -164,18 +191,31 @@ class LineSearchBase(Module, ABC):
             derivative = - sum(t.sum() for t in torch._foreach_mul([p.grad if p.grad is not None
                                                                     else torch.zeros_like(p) for p in params], update))
 
-        return loss, tofloat(derivative)
+        assert var.grad is not None
+        return loss, tofloat(derivative), var.grad
 
-    def evaluate_step_size(self, step_size: float, var: Var, backward:bool=False):
+    def _loss_derivative(self, step_size: float, var: Var, closure,
+                         params: list[torch.Tensor], update: list[torch.Tensor]):
+        return self._loss_derivative_gradient(step_size=step_size, var=var,closure=closure,params=params,update=update)[:2]
+
+    def evaluate_f(self, step_size: float, var: Var, backward:bool=False):
+        """evaluate function value at alpha `step_size`."""
         closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
         return self._loss(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_update(),backward=backward)
 
-    def evaluate_step_size_loss_and_derivative(self, step_size: float, var: Var):
+    def evaluate_f_d(self, step_size: float, var: Var):
+        """evaluate function value and directional derivative in the direction of the update at step size `step_size`."""
         closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
         return self._loss_derivative(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_update())
 
+    def evaluate_f_d_g(self, step_size: float, var: Var):
+        """evaluate function value, directional derivative, and gradient list at step size `step_size`."""
+        closure = var.closure
+        if closure is None: raise RuntimeError('line search requires closure')
+        return self._loss_derivative_gradient(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_update())
+
     def make_objective(self, var: Var, backward:bool=False):
         closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
@@ -186,6 +226,11 @@ class LineSearchBase(Module, ABC):
         if closure is None: raise RuntimeError('line search requires closure')
         return partial(self._loss_derivative, var=var, closure=closure, params=var.params, update=var.get_update())
 
+    def make_objective_with_derivative_and_gradient(self, var: Var):
+        closure = var.closure
+        if closure is None: raise RuntimeError('line search requires closure')
+        return partial(self._loss_derivative_gradient, var=var, closure=closure, params=var.params, update=var.get_update())
+
     @abstractmethod
     def search(self, update: list[torch.Tensor], var: Var) -> float:
         """Finds the step size to use"""
@@ -193,7 +238,9 @@ class LineSearchBase(Module, ABC):
     @torch.no_grad
     def step(self, var: Var) -> Var:
         self._reset()
+
         params = var.params
+        self._initial_params = [p.clone() for p in params]
         update = var.get_update()
 
         try:
@@ -206,7 +253,6 @@ class LineSearchBase(Module, ABC):
 
         # this is last module - set step size to found step_size times lr
         if var.is_last:
-
             if var.last_module_lrs is None:
                 self.set_step_size_(step_size, params=params, update=update)
 
@@ -223,17 +269,62 @@ class LineSearchBase(Module, ABC):
 
 
 
-# class GridLineSearch(LineSearch):
-#     """Mostly for testing, this is not practical"""
-#     def __init__(self, start, end, num):
-#         defaults = dict(start=start,end=end,num=num)
-#         super().__init__(defaults)
-
-#     @torch.no_grad
-#     def search(self, update, var):
-#         start,end,num=itemgetter('start','end','num')(self.settings[var.params[0]])
-
-#         for lr in torch.linspace(start,end,num):
-#             self.evaluate_step_size(lr.item(), var=var, backward=False)
+class GridLineSearch(LineSearchBase):
+    """"""
+    def __init__(self, start, end, num):
+        defaults = dict(start=start,end=end,num=num)
+        super().__init__(defaults)
 
-#         return self._best_step_size
+    @torch.no_grad
+    def search(self, update, var):
+        start,end,num=itemgetter('start','end','num')(self.defaults)
+
+        for lr in torch.linspace(start,end,num):
+            self.evaluate_f(lr.item(), var=var, backward=False)
+
+        return self._best_step_size
+
+
+def sufficient_decrease(f_0, g_0, f_a, a, c):
+    return f_a < f_0 + c*a*min(g_0, 0)
+
+def curvature(g_0, g_a, c):
+    if g_0 > 0: return True
+    return g_a >= c * g_0
+
+def strong_curvature(g_0, g_a, c):
+    """same as curvature condition except curvature can't be too positive (which indicates overstep)"""
+    if g_0 > 0: return True
+    return abs(g_a) <= c * abs(g_0)
+
+def wolfe(f_0, g_0, f_a, g_a, a, c1, c2):
+    return sufficient_decrease(f_0, g_0, f_a, a, c1) and curvature(g_0, g_a, c2)
+
+def strong_wolfe(f_0, g_0, f_a, g_a, a, c1, c2):
+    return sufficient_decrease(f_0, g_0, f_a, a, c1) and strong_curvature(g_0, g_a, c2)
+
+def goldstein(f_0, g_0, f_a, a, c):
+    """same as armijo (sufficient_decrease) but additional lower bound"""
+    g_0 = min(g_0, 0)
+    return f_0 + (1-c)*a*g_0 < f_a < f_0 + c*a*g_0
+
+TerminationCondition = Literal["armijo", "curvature", "strong_curvature", "wolfe", "strong_wolfe", "goldstein", "decrease"]
+def termination_condition(
+    condition: TerminationCondition,
+    f_0,
+    g_0,
+    f_a,
+    g_a: Any | None,
+    a,
+    c,
+    c2=None,
+):
+    if not math.isfinite(f_a): return False
+    if condition == 'armijo': return sufficient_decrease(f_0, g_0, f_a, a, c)
+    if condition == 'curvature': return curvature(g_0, g_a, c)
+    if condition == 'strong_curvature': return strong_curvature(g_0, g_a, c)
+    if condition == 'wolfe': return wolfe(f_0, g_0, f_a, g_a, a, c, c2)
+    if condition == 'strong_wolfe': return strong_wolfe(f_0, g_0, f_a, g_a, a, c, c2)
+    if condition == 'goldstein': return goldstein(f_0, g_0, f_a, a, c)
+    if condition == 'decrease': return f_a < f_0
+    raise ValueError(f"unknown condition {condition}")

torchzero/modules/line_search/scipy.py

@@ -1,3 +1,4 @@
+import math
 from collections.abc import Mapping
 from operator import itemgetter
 
@@ -17,6 +18,7 @@ class ScipyMinimizeScalar(LineSearchBase):
         bounds (Sequence | None, optional):
             For method ‘bounded’, bounds is mandatory and must have two finite items corresponding to the optimization bounds. Defaults to None.
         tol (float | None, optional): Tolerance for termination. Defaults to None.
+        prev_init (bool, optional): uses previous step size as initial guess for the line search.
         options (dict | None, optional): A dictionary of solver options. Defaults to None.
 
     For more details on methods and arguments refer to https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize_scalar.html
@@ -29,9 +31,10 @@ class ScipyMinimizeScalar(LineSearchBase):
         bracket=None,
         bounds=None,
         tol: float | None = None,
+        prev_init: bool = False,
         options=None,
     ):
-        defaults = dict(method=method,bracket=bracket,bounds=bounds,tol=tol,options=options,maxiter=maxiter)
+        defaults = dict(method=method,bracket=bracket,bounds=bounds,tol=tol,options=options,maxiter=maxiter, prev_init=prev_init)
         super().__init__(defaults)
 
         import scipy.optimize
@@ -42,11 +45,20 @@ class ScipyMinimizeScalar(LineSearchBase):
     def search(self, update, var):
         objective = self.make_objective(var=var)
         method, bracket, bounds, tol, options, maxiter = itemgetter(
-            'method', 'bracket', 'bounds', 'tol', 'options', 'maxiter')(self.settings[var.params[0]])
+            'method', 'bracket', 'bounds', 'tol', 'options', 'maxiter')(self.defaults)
 
         if maxiter is not None:
             options = dict(options) if isinstance(options, Mapping) else {}
             options['maxiter'] = maxiter
 
-        res = self.scopt.minimize_scalar(objective, method=method, bracket=bracket, bounds=bounds, tol=tol, options=options)
-        return res.x
+        if self.defaults["prev_init"] and "x_prev" in self.global_state:
+            if bracket is None: bracket = (0, 1)
+            bracket = (*bracket[:-1], self.global_state["x_prev"])
+
+        x = self.scopt.minimize_scalar(objective, method=method, bracket=bracket, bounds=bounds, tol=tol, options=options).x # pyright:ignore[reportAttributeAccessIssue]
+
+        max = torch.finfo(var.params[0].dtype).max / 2
+        if (not math.isfinite(x)) or abs(x) >= max: x = 0
+
+        self.global_state['x_prev'] = x
+        return x