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

torchzero/modules/second_order/newton_cg.py

@@ -1,11 +1,14 @@
-from typing import Literal, overload
+import warnings
+import math
+from typing import Literal, cast
+from operator import itemgetter
 import torch
 
-from ...utils import TensorList, as_tensorlist, NumberList
+from ...core import Chainable, Module, apply_transform
+from ...utils import TensorList, as_tensorlist, tofloat
 from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
-
-from ...core import Chainable, apply_transform, Module
-from ...utils.linalg.solve import cg, steihaug_toint_cg, minres
+from ...utils.linalg.solve import cg, minres, find_within_trust_radius
+from ..trust_region.trust_region import default_radius
 
 class NewtonCG(Module):
     """Newton's method with a matrix-free conjugate gradient or minimial-residual solver.
@@ -88,20 +91,25 @@ class NewtonCG(Module):
     def __init__(
         self,
         maxiter: int | None = None,
-        tol: float = 1e-4,
+        tol: float = 1e-8,
         reg: float = 1e-8,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
         solver: Literal['cg', 'minres', 'minres_npc'] = 'cg',
         h: float = 1e-3,
+        miniter:int = 1,
         warm_start=False,
         inner: Chainable | None = None,
     ):
-        defaults = dict(tol=tol, maxiter=maxiter, reg=reg, hvp_method=hvp_method, solver=solver, h=h, warm_start=warm_start)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner']
         super().__init__(defaults,)
 
         if inner is not None:
             self.set_child('inner', inner)
 
+        self._num_hvps = 0
+        self._num_hvps_last_step = 0
+
     @torch.no_grad
     def step(self, var):
         params = TensorList(var.params)
@@ -117,11 +125,13 @@ class NewtonCG(Module):
         h = settings['h']
         warm_start = settings['warm_start']
 
+        self._num_hvps_last_step = 0
         # ---------------------- Hessian vector product function --------------------- #
         if hvp_method == 'autograd':
             grad = var.get_grad(create_graph=True)
 
             def H_mm(x):
+                self._num_hvps_last_step += 1
                 with torch.enable_grad():
                     return TensorList(hvp(params, grad, x, retain_graph=True))
 
@@ -132,10 +142,12 @@ class NewtonCG(Module):
 
             if hvp_method == 'forward':
                 def H_mm(x):
+                    self._num_hvps_last_step += 1
                     return TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1])
 
             elif hvp_method == 'central':
                 def H_mm(x):
+                    self._num_hvps_last_step += 1
                     return TensorList(hvp_fd_central(closure, params, x, h=h, normalize=True)[1])
 
             else:
@@ -153,26 +165,28 @@ class NewtonCG(Module):
         if warm_start: x0 = self.get_state(params, 'prev_x', cls=TensorList) # initialized to 0 which is default anyway
 
         if solver == 'cg':
-            x = cg(A_mm=H_mm, b=b, x0_=x0, tol=tol, maxiter=maxiter, reg=reg)
+            d, _ = cg(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, miniter=self.defaults["miniter"],reg=reg)
 
         elif solver == 'minres':
-            x = minres(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=False)
+            d = minres(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=False)
 
         elif solver == 'minres_npc':
-            x = minres(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=True)
+            d = minres(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=True)
 
         else:
             raise ValueError(f"Unknown solver {solver}")
 
         if warm_start:
             assert x0 is not None
-            x0.copy_(x)
+            x0.copy_(d)
+
+        var.update = d
 
-        var.update = x
+        self._num_hvps += self._num_hvps_last_step
         return var
 
 
-class TruncatedNewtonCG(Module):
+class NewtonCGSteihaug(Module):
     """Trust region Newton's method with a matrix-free Steihaug-Toint conjugate gradient or MINRES solver.
 
     This optimizer implements Newton's method using a matrix-free conjugate
@@ -245,49 +259,61 @@ class TruncatedNewtonCG(Module):
     def __init__(
         self,
         maxiter: int | None = None,
-        eta: float= 1e-6,
-        nplus: float = 2,
+        eta: float= 0.0,
+        nplus: float = 3.5,
         nminus: float = 0.25,
+        rho_good: float = 0.99,
+        rho_bad: float = 1e-4,
         init: float = 1,
-        tol: float = 1e-4,
+        tol: float = 1e-8,
         reg: float = 1e-8,
-        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
-        solver: Literal['cg', 'minres', 'minres_npc'] = 'cg',
+        hvp_method: Literal["forward", "central"] = "forward",
+        solver: Literal['cg', "minres"] = 'cg',
         h: float = 1e-3,
-        max_attempts: int = 10,
+        max_attempts: int = 100,
+        max_history: int = 100,
+        boundary_tol: float = 1e-1,
+        miniter: int = 1,
+        rms_beta: float | None = None,
+        adapt_tol: bool = True,
+        npc_terminate: bool = False,
         inner: Chainable | None = None,
     ):
-        defaults = dict(tol=tol, maxiter=maxiter, reg=reg, hvp_method=hvp_method, h=h, eta=eta, nplus=nplus, nminus=nminus, init=init, max_attempts=max_attempts, solver=solver)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner']
         super().__init__(defaults,)
 
         if inner is not None:
             self.set_child('inner', inner)
 
+        self._num_hvps = 0
+        self._num_hvps_last_step = 0
+
     @torch.no_grad
     def step(self, var):
         params = TensorList(var.params)
         closure = var.closure
         if closure is None: raise RuntimeError('NewtonCG requires closure')
 
-        settings = self.settings[params[0]]
-        tol = settings['tol']
-        reg = settings['reg']
-        maxiter = settings['maxiter']
-        hvp_method = settings['hvp_method']
-        h = settings['h']
-        max_attempts = settings['max_attempts']
-        solver = settings['solver'].lower().strip()
+        tol = self.defaults['tol'] * self.global_state.get('tol_mul', 1)
+        solver = self.defaults['solver'].lower().strip()
 
-        eta = settings['eta']
-        nplus = settings['nplus']
-        nminus = settings['nminus']
-        init = settings['init']
+        (reg, maxiter, hvp_method, h, max_attempts, boundary_tol,
+         eta, nplus, nminus, rho_good, rho_bad, init, npc_terminate,
+         miniter, max_history, adapt_tol) = itemgetter(
+             "reg", "maxiter", "hvp_method", "h", "max_attempts", "boundary_tol",
+             "eta", "nplus", "nminus", "rho_good", "rho_bad", "init", "npc_terminate",
+             "miniter", "max_history", "adapt_tol",
+        )(self.defaults)
+
+        self._num_hvps_last_step = 0
 
         # ---------------------- Hessian vector product function --------------------- #
         if hvp_method == 'autograd':
             grad = var.get_grad(create_graph=True)
 
             def H_mm(x):
+                self._num_hvps_last_step += 1
                 with torch.enable_grad():
                     return TensorList(hvp(params, grad, x, retain_graph=True))
 
@@ -298,77 +324,112 @@ class TruncatedNewtonCG(Module):
 
             if hvp_method == 'forward':
                 def H_mm(x):
+                    self._num_hvps_last_step += 1
                     return TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1])
 
             elif hvp_method == 'central':
                 def H_mm(x):
+                    self._num_hvps_last_step += 1
                     return TensorList(hvp_fd_central(closure, params, x, h=h, normalize=True)[1])
 
             else:
                 raise ValueError(hvp_method)
 
 
-        # -------------------------------- inner step -------------------------------- #
+        # ------------------------- update RMS preconditioner ------------------------ #
         b = var.get_update()
+        P_mm = None
+        rms_beta = self.defaults["rms_beta"]
+        if rms_beta is not None:
+            exp_avg_sq = self.get_state(params, "exp_avg_sq", init=b, cls=TensorList)
+            exp_avg_sq.mul_(rms_beta).addcmul(b, b, value=1-rms_beta)
+            exp_avg_sq_sqrt = exp_avg_sq.sqrt().add_(1e-8)
+            def _P_mm(x):
+                return x / exp_avg_sq_sqrt
+            P_mm = _P_mm
+
+        # -------------------------------- inner step -------------------------------- #
         if 'inner' in self.children:
             b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
         b = as_tensorlist(b)
 
-        # ---------------------------------- run cg ---------------------------------- #
+        # ------------------------------- trust region ------------------------------- #
         success = False
-        x = None
+        d = None
+        x0 = [p.clone() for p in params]
+        solution = None
+
         while not success:
             max_attempts -= 1
             if max_attempts < 0: break
 
-            trust_region = self.global_state.get('trust_region', init)
-            if trust_region < 1e-8 or trust_region > 1e8:
-                trust_region = self.global_state['trust_region'] = init
-
-            if solver == 'cg':
-                x = steihaug_toint_cg(A_mm=H_mm, b=b, trust_region=trust_region, tol=tol, maxiter=maxiter, reg=reg)
-
-            elif solver == 'minres':
-                x = minres(A_mm=H_mm, b=b, trust_region=trust_region, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=False)
-
-            elif solver == 'minres_npc':
-                x = minres(A_mm=H_mm, b=b, trust_region=trust_region, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=True)
-
-            else:
-                raise ValueError(f"unknown solver {solver}")
-
-            # ------------------------------- trust region ------------------------------- #
-            Hx = H_mm(x)
-            pred_reduction = b.dot(x) - 0.5 * x.dot(Hx)
-
-            params -= x
-            loss_star = closure(False)
-            params += x
-            reduction = var.get_loss(False) - loss_star
-
-            rho = reduction / (pred_reduction.clip(min=1e-8))
-
-            # failed step
-            if rho < 0.25:
-                self.global_state['trust_region'] = trust_region * nminus
-
-            # very good step
-            elif rho > 0.75:
-                diff = trust_region - x.abs()
-                if (diff.global_min() / trust_region) > 1e-4: # hits boundary
-                    self.global_state['trust_region'] = trust_region * nplus
-
-            # if the ratio is high enough then accept the proposed step
-            if rho > eta:
-                success = True
+            trust_radius = self.global_state.get('trust_radius', init)
+
+            # -------------- make sure trust radius isn't too small or large ------------- #
+            finfo = torch.finfo(x0[0].dtype)
+            if trust_radius < finfo.tiny * 2:
+                trust_radius = self.global_state['trust_radius'] = init
+                if adapt_tol:
+                    self.global_state["tol_mul"] = self.global_state.get("tol_mul", 1) * 0.1
+
+            elif trust_radius > finfo.max / 2:
+                trust_radius = self.global_state['trust_radius'] = init
+
+            # ----------------------------------- solve ---------------------------------- #
+            d = None
+            if solution is not None and solution.history is not None:
+                d = find_within_trust_radius(solution.history, trust_radius)
+
+            if d is None:
+                if solver == 'cg':
+                    d, solution = cg(
+                        A_mm=H_mm,
+                        b=b,
+                        tol=tol,
+                        maxiter=maxiter,
+                        reg=reg,
+                        trust_radius=trust_radius,
+                        miniter=miniter,
+                        npc_terminate=npc_terminate,
+                        history_size=max_history,
+                        P_mm=P_mm,
+                    )
+
+                elif solver == 'minres':
+                    d = minres(A_mm=H_mm, b=b, trust_radius=trust_radius, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=npc_terminate)
+
+                else:
+                    raise ValueError(f"unknown solver {solver}")
+
+            # ---------------------------- update trust radius --------------------------- #
+            self.global_state["trust_radius"], success = default_radius(
+                params=params,
+                closure=closure,
+                f=tofloat(var.get_loss(False)),
+                g=b,
+                H=H_mm,
+                d=d,
+                trust_radius=trust_radius,
+                eta=eta,
+                nplus=nplus,
+                nminus=nminus,
+                rho_good=rho_good,
+                rho_bad=rho_bad,
+                boundary_tol=boundary_tol,
+
+                init=init, # init isn't used because check_overflow=False
+                state=self.global_state, # not used
+                settings=self.defaults, # not used
+                check_overflow=False, # this is checked manually to adapt tolerance
+            )
 
-        assert x is not None
+        # --------------------------- assign new direction --------------------------- #
+        assert d is not None
         if success:
-            var.update = x
+            var.update = d
 
         else:
             var.update = params.zeros_like()
 
-        return var
-
-
+        self._num_hvps += self._num_hvps_last_step
+        return var

torchzero/modules/smoothing/__init__.py

@@ -1,2 +1,2 @@
 from .laplacian import LaplacianSmoothing
-from .gaussian import GaussianHomotopy
+from .sampling import GradientSampling

torchzero/modules/smoothing/sampling.py

@@ -0,0 +1,300 @@
+import math
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from contextlib import nullcontext
+from functools import partial
+from typing import Literal, cast
+
+import torch
+
+from ...core import Chainable, Modular, Module, Var
+from ...core.reformulation import Reformulation
+from ...utils import Distributions, NumberList, TensorList
+from ..termination import TerminationCriteriaBase, make_termination_criteria
+
+
+def _reset_except_self(optimizer: Modular, var: Var, self: Module):
+    for m in optimizer.unrolled_modules:
+        if m is not self:
+            m.reset()
+
+
+class GradientSampling(Reformulation):
+    """Samples and aggregates gradients and values at perturbed points.
+
+    This module can be used for gaussian homotopy and gradient sampling methods.
+
+    Args:
+        modules (Chainable | None, optional):
+            modules that will be optimizing the modified objective.
+            if None, returns gradient of the modified objective as the update. Defaults to None.
+        sigma (float, optional): initial magnitude of the perturbations. Defaults to 1.
+        n (int, optional): number of perturbations per step. Defaults to 100.
+        aggregate (str, optional):
+            how to aggregate values and gradients
+            - "mean" - uses mean of the gradients, as in gaussian homotopy.
+            - "max" - uses element-wise maximum of the gradients.
+            - "min" - uses element-wise minimum of the gradients.
+            - "min-norm" - picks gradient with the lowest norm.
+
+            Defaults to 'mean'.
+        distribution (Distributions, optional): distribution for random perturbations. Defaults to 'gaussian'.
+        include_x0 (bool, optional): whether to include gradient at un-perturbed point. Defaults to True.
+        fixed (bool, optional):
+            if True, perturbations do not get replaced by new random perturbations until termination criteria is satisfied. Defaults to True.
+        pre_generate (bool, optional):
+            if True, perturbations are pre-generated before each step.
+            This requires more memory to store all of them,
+            but ensures they do not change when closure is evaluated multiple times.
+            Defaults to True.
+        termination (TerminationCriteriaBase | Sequence[TerminationCriteriaBase] | None, optional):
+            a termination criteria module, sigma will be multiplied by ``decay`` when termination criteria is satisfied,
+            and new perturbations will be generated if ``fixed``. Defaults to None.
+        decay (float, optional): sigma multiplier on termination criteria. Defaults to 2/3.
+        reset_on_termination (bool, optional): whether to reset states of all other modules on termination. Defaults to True.
+        sigma_strategy (str | None, optional):
+            strategy for adapting sigma. If condition is satisfied, sigma is multiplied by ``sigma_nplus``,
+            otherwise it is multiplied by ``sigma_nminus``.
+            - "grad-norm" - at least ``sigma_target`` gradients should have lower norm than at un-perturbed point.
+            - "value" - at least ``sigma_target`` values (losses) should be lower than at un-perturbed point.
+            - None - doesn't use adaptive sigma.
+
+            This introduces a side-effect to the closure, so it should be left at None of you use
+            trust region or line search to optimize the modified objective.
+            Defaults to None.
+        sigma_target (int, optional):
+            number of elements to satisfy the condition in ``sigma_strategy``. Defaults to 1.
+        sigma_nplus (float, optional): sigma multiplier when ``sigma_strategy`` condition is satisfied. Defaults to 4/3.
+        sigma_nminus (float, optional): sigma multiplier when ``sigma_strategy`` condition is not satisfied. Defaults to 2/3.
+        seed (int | None, optional): seed. Defaults to None.
+    """
+    def __init__(
+        self,
+        modules: Chainable | None = None,
+        sigma: float = 1.,
+        n:int = 100,
+        aggregate: Literal['mean', 'max', 'min', 'min-norm', 'min-value'] = 'mean',
+        distribution: Distributions = 'gaussian',
+        include_x0: bool = True,
+
+        fixed: bool=True,
+        pre_generate: bool = True,
+        termination: TerminationCriteriaBase | Sequence[TerminationCriteriaBase] | None = None,
+        decay: float = 2/3,
+        reset_on_termination: bool = True,
+
+        sigma_strategy: Literal['grad-norm', 'value'] | None = None,
+        sigma_target: int | float = 0.2,
+        sigma_nplus: float = 4/3,
+        sigma_nminus: float = 2/3,
+
+        seed: int | None = None,
+    ):
+
+        defaults = dict(sigma=sigma, n=n, aggregate=aggregate, distribution=distribution, seed=seed, include_x0=include_x0, fixed=fixed, decay=decay, reset_on_termination=reset_on_termination, sigma_strategy=sigma_strategy, sigma_target=sigma_target, sigma_nplus=sigma_nplus, sigma_nminus=sigma_nminus, pre_generate=pre_generate)
+        super().__init__(defaults, modules)
+
+        if termination is not None:
+            self.set_child('termination', make_termination_criteria(extra=termination))
+
+    @torch.no_grad
+    def pre_step(self, var):
+        params = TensorList(var.params)
+
+        fixed = self.defaults['fixed']
+
+        # check termination criteria
+        if 'termination' in self.children:
+            termination = cast(TerminationCriteriaBase, self.children['termination'])
+            if termination.should_terminate(var):
+
+                # decay sigmas
+                states = [self.state[p] for p in params]
+                settings = [self.settings[p] for p in params]
+
+                for state, setting in zip(states, settings):
+                    if 'sigma' not in state: state['sigma'] = setting['sigma']
+                    state['sigma'] *= setting['decay']
+
+                # reset on sigmas decay
+                if self.defaults['reset_on_termination']:
+                    var.post_step_hooks.append(partial(_reset_except_self, self=self))
+
+                # clear perturbations
+                self.global_state.pop('perts', None)
+
+        # pre-generate perturbations if not already pre-generated or not fixed
+        if self.defaults['pre_generate'] and (('perts' not in self.global_state) or (not fixed)):
+            states = [self.state[p] for p in params]
+            settings = [self.settings[p] for p in params]
+
+            n = self.defaults['n'] - self.defaults['include_x0']
+            generator = self.get_generator(params[0].device, self.defaults['seed'])
+
+            perts = [params.sample_like(self.defaults['distribution'], generator=generator) for _ in range(n)]
+
+            self.global_state['perts'] = perts
+
+    @torch.no_grad
+    def closure(self, backward, closure, params, var):
+        params = TensorList(params)
+        loss_agg = None
+        grad_agg = None
+
+        states = [self.state[p] for p in params]
+        settings = [self.settings[p] for p in params]
+        sigma_inits = [s['sigma'] for s in settings]
+        sigmas = [s.setdefault('sigma', si) for s, si in zip(states, sigma_inits)]
+
+        include_x0 = self.defaults['include_x0']
+        pre_generate = self.defaults['pre_generate']
+        aggregate: Literal['mean', 'max', 'min', 'min-norm', 'min-value'] = self.defaults['aggregate']
+        sigma_strategy: Literal['grad-norm', 'value'] | None = self.defaults['sigma_strategy']
+        distribution = self.defaults['distribution']
+        generator = self.get_generator(params[0].device, self.defaults['seed'])
+
+
+        n_finite = 0
+        n_good = 0
+        f_0 = None; g_0 = None
+
+        # evaluate at x_0
+        if include_x0:
+            f_0 = cast(torch.Tensor, var.get_loss(backward=backward))
+
+            isfinite = math.isfinite(f_0)
+            if isfinite:
+                n_finite += 1
+                loss_agg = f_0
+
+            if backward:
+                g_0 = var.get_grad()
+                if isfinite: grad_agg = g_0
+
+        # evaluate at x_0 + p for each perturbation
+        if pre_generate:
+            perts = self.global_state['perts']
+        else:
+            perts = [None] * (self.defaults['n'] - include_x0)
+
+        x_0 = [p.clone() for p in params]
+
+        for pert in perts:
+            loss = None; grad = None
+
+            # generate if not pre-generated
+            if pert is None:
+                pert = params.sample_like(distribution, generator=generator)
+
+            # add perturbation and evaluate
+            pert = pert * sigmas
+            torch._foreach_add_(params, pert)
+
+            with torch.enable_grad() if backward else nullcontext():
+                loss = closure(backward)
+
+            if math.isfinite(loss):
+                n_finite += 1
+
+                # add loss
+                if loss_agg is None:
+                    loss_agg = loss
+                else:
+                    if aggregate == 'mean':
+                        loss_agg += loss
+
+                    elif (aggregate=='min') or (aggregate=='min-value') or (aggregate=='min-norm' and not backward):
+                        loss_agg = loss_agg.clamp(max=loss)
+
+                    elif aggregate == 'max':
+                        loss_agg = loss_agg.clamp(min=loss)
+
+                # add grad
+                if backward:
+                    grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+                    if grad_agg is None:
+                        grad_agg = grad
+                    else:
+                        if aggregate == 'mean':
+                            torch._foreach_add_(grad_agg, grad)
+
+                        elif aggregate == 'min':
+                            grad_agg_abs = torch._foreach_abs(grad_agg)
+                            torch._foreach_minimum_(grad_agg_abs, torch._foreach_abs(grad))
+                            grad_agg = [g_abs.copysign(g) for g_abs, g in zip(grad_agg_abs, grad_agg)]
+
+                        elif aggregate == 'max':
+                            grad_agg_abs = torch._foreach_abs(grad_agg)
+                            torch._foreach_maximum_(grad_agg_abs, torch._foreach_abs(grad))
+                            grad_agg = [g_abs.copysign(g) for g_abs, g in zip(grad_agg_abs, grad_agg)]
+
+                        elif aggregate == 'min-norm':
+                            if TensorList(grad).global_vector_norm() < TensorList(grad_agg).global_vector_norm():
+                                grad_agg = grad
+                                loss_agg = loss
+
+                        elif aggregate == 'min-value':
+                            if loss < loss_agg:
+                                grad_agg = grad
+                                loss_agg = loss
+
+            # undo perturbation
+            torch._foreach_copy_(params, x_0)
+
+            # adaptive sigma
+            # by value
+            if sigma_strategy == 'value':
+                if f_0 is None:
+                    with torch.enable_grad() if backward else nullcontext():
+                        f_0 = closure(False)
+
+                if loss < f_0:
+                    n_good += 1
+
+            # by gradient norm
+            elif sigma_strategy == 'grad-norm' and backward and math.isfinite(loss):
+                assert grad is not None
+                if g_0 is None:
+                    with torch.enable_grad() if backward else nullcontext():
+                        closure()
+                        g_0 = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+
+                if TensorList(grad).global_vector_norm() < TensorList(g_0).global_vector_norm():
+                    n_good += 1
+
+        # update sigma if strategy is enabled
+        if sigma_strategy is not None:
+
+            sigma_target = self.defaults['sigma_target']
+            if isinstance(sigma_target, float):
+                sigma_target = int(max(1, n_finite * sigma_target))
+
+            if n_good >= sigma_target:
+                key = 'sigma_nplus'
+            else:
+                key = 'sigma_nminus'
+
+            for p in params:
+                self.state[p]['sigma'] *= self.settings[p][key]
+
+        # if no finite losses, just return inf
+        if n_finite == 0:
+            assert loss_agg is None and grad_agg is None
+            loss = torch.tensor(torch.inf, dtype=params[0].dtype, device=params[0].device)
+            grad = [torch.full_like(p, torch.inf) for p in params]
+            return loss, grad
+
+        assert loss_agg is not None
+
+        # no post processing needed when aggregate is 'max', 'min', 'min-norm', 'min-value'
+        if aggregate != 'mean':
+            return loss_agg, grad_agg
+
+        # on mean divide by number of evals
+        loss_agg /= n_finite
+
+        if backward:
+            assert grad_agg is not None
+            torch._foreach_div_(grad_agg, n_finite)
+
+        return loss_agg, grad_agg

torchzero/modules/step_size/__init__.py

@@ -1,2 +1,2 @@
 from .lr import LR, StepSize, Warmup, WarmupNormClip, RandomStepSize
-from .adaptive import PolyakStepSize, BarzilaiBorwein
+from .adaptive import PolyakStepSize, BarzilaiBorwein, BBStab, AdGD