torchzero 0.0.1__py3-none-any.whl

torchzero/modules/scheduling/lr_schedulers.py

@@ -0,0 +1,131 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, overload, TYPE_CHECKING
+import random
+
+import torch
+from ...core import OptimizerModule
+
+
+if TYPE_CHECKING:
+    from ...optim import Modular
+
+
+# LR SCHEDULING MOVED TO LR MODULE
+
+# def _set_momentum_hook(optimizer, state, momentum):
+#     for module in optimizer.unrolled_modules:
+#         if 'momentum' in module.defaults:
+#             for g in module.param_groups:
+#                 g['momentum'] = momentum
+#         if 'beta1' in module.defaults:
+#             for g in module.param_groups:
+#                 g['beta1'] = momentum
+
+# def _add_scheduler_hook(opt: "Modular", scheduler_cls, id):
+#     """post-init hook that sets `scheduler_step_fn` to the scheduler step."""
+#     # get LR module
+#     lr_module = opt.get_lr_module()
+
+#     # get current LRScheduler module
+#     scheds = [i for i in opt.unrolled_modules if isinstance(i, LRScheduler)]
+#     scheds = [i for i in scheds if i.id == id]
+#     if len(scheds) != 1:
+#         raise RuntimeError(f"more than 1 module with id {id}: {scheds}")
+
+#     sch_module = scheds[0]
+
+#     # make a scheduler and save the step function
+#     scheduler = scheduler_cls(lr_module)
+#     sch_module.scheduler_step_fn = scheduler.step
+
+
+# class LRScheduler(OptimizerModule):
+#     """Use any pytorch lr scheduler.
+
+#     Important - the lr is applied multiplicatively and multiplies with learning rate of other modules,
+#     so usually base learning rate of the lr scheduler, such as `max_lr` for OneCycleLR, should be set to 1.
+
+#     Args:
+#         lr_scheduler (Callable[[torch.optim.Optimizer], torch.optim.lr_scheduler.LRScheduler  |  Any]):
+#             something like:
+#             .. code:: py
+#             lambda opt: OneCycleLR(opt, max_lr = 1, total_steps = 60000)
+#         update_every (int, optional):
+#             call `step` every n steps, useful for schedulers that only step once per epoch. Defaults to 1.
+#         cycle_momentum (bool, optional):
+#             enables support for cycling momentum with schedulers that support it, such as `OneCycleLR`.
+#             Unlike lr, momentum is not applied multiplicatively, but set to all other modules with
+#             `momentum` or `beta` settings. Has no effect if there are no modules that support momentum. Defaults to False.
+#         init_lr (float, optional):
+#             initial lr, I believe most lr schedulers ignore this. Defaults to 1.
+#         init_momentum (float, optional):
+#             initial init_momentum, I believe most lr schedulers ignore this.
+#             Has no effect if `cycle_momentum` is False or there are no modules that support momentum. Defaults to 0.
+#     """
+#     def __init__(
+#         self,
+#         lr_scheduler: Callable[[torch.optim.Optimizer], torch.optim.lr_scheduler.LRScheduler | Any],
+#         step_every: int = 1,
+#         cycle_momentum: bool = True,
+#     ):
+#         super().__init__({})
+#         scheduler = lr_scheduler(self.dummy_opt)
+#         self.update_every = step_every
+#         self.cycle_momentum = cycle_momentum
+
+#         self.scheduler_step_fn = scheduler.step
+#         self.cur = 0
+#         self.cur_lr = init_lr
+#         self.cur_momentum = init_momentum
+
+#         self.id = random.random()
+
+#     def step(self, state):
+#         if self.cur % self.update_every == 0:
+#             self.scheduler_step_fn()
+#             self.cur_lr = self.dummy_opt.first_param_group['lr']
+#             self.cur_momentum = self.dummy_opt.first_param_group['momentum']
+
+#         params = self.get_params()
+#         ascent = state.maybe_use_grad_(params)
+#         ascent *= self.cur_lr
+
+#         if self.cycle_momentum:
+#             state.add_post_step_hook(partial(_set_momentum_hook, momentum = self.cur_momentum))
+
+class LRWarmup(OptimizerModule):
+    """Linear learning rate warmup.
+
+    Args:
+        n_steps (int): number of warmup steps.
+        start_lr (float, optional): initial lr. Defaults to 1e-8.
+        end_lr (float, optional): final lr. Defaults to 1.
+        delay_steps (int, optional): number of `start_lr` steps before starting the warmup. Defaults to 0.
+    """
+    def __init__(self, n_steps: int, start_lr: float = 1e-8, end_lr: float = 1, delay_steps: int = 0):
+
+        super().__init__({})
+        self.n_steps = n_steps
+        self.start_lr = start_lr
+        self.end_lr = end_lr
+        self.delay_steps = delay_steps
+
+        self.cur = 0
+
+    def _update(self, state, ascent):
+        if self.cur < self.delay_steps:
+            if self.start_lr != 1: ascent *= self.start_lr
+
+        elif self.cur >= self.n_steps + self.delay_steps:
+            if self.end_lr != 1: ascent *= self.end_lr
+
+        else:
+            remaining = (self.n_steps - (self.cur-self.delay_steps)) / self.n_steps
+            lr = (self.start_lr * remaining) + self.end_lr * (1 - remaining)
+            ascent *= lr
+
+        self.cur += 1
+        return ascent
+
+

torchzero/modules/scheduling/step_size.py

@@ -0,0 +1,80 @@
+import random
+from typing import Any
+
+from ...core import OptimizerModule
+from ...tensorlist import TensorList
+
+
+class PolyakStepSize(OptimizerModule):
+    """Polyak step-size. Meant to be used at the beginning when ascent is the gradient but other placements may work.
+    This can also work with SGD as SPS (Stochastic Polyak Step-Size) seems to use the same formula.
+
+    Args:
+        max (float | None, optional): maximum possible step size. Defaults to None.
+        min_obj_value (int, optional): (estimated) minimal possible value of the objective function (lowest possible loss). Defaults to 0.
+        use_grad (bool, optional):
+            if True, uses dot product of update and gradient to compute the step size.
+            Otherwise, dot product of update with itself is used, which has no geometric meaning so it probably won't work well.
+            Defaults to True.
+        parameterwise (bool, optional):
+            if True, calculate Polyak step-size for each parameter separately,
+            if False calculate one global step size for all parameters. Defaults to False.
+        alpha (float, optional): multiplier to Polyak step-size. Defaults to 1.
+    """
+    def __init__(self, max: float | None = None, min_obj_value: float = 0, use_grad=True, parameterwise=False, alpha: float = 1):
+
+        defaults = dict(alpha = alpha)
+        super().__init__(defaults)
+        self.max = max
+        self.min_obj_value = min_obj_value
+        self.use_grad = use_grad
+        self.parameterwise = parameterwise
+
+    def _update(self, state, ascent):
+        if state.closure is None: raise ValueError("PolyakStepSize requires closure")
+        if state.fx0 is None: state.fx0 = state.closure(False) # can only happen when placed after SPSA
+
+        alpha = self.get_group_key('alpha')
+
+        if self.parameterwise:
+            if self.use_grad: denom = (ascent*state.maybe_compute_grad_(self.get_params())).mean()
+            else: denom = ascent.pow(2).mean()
+            polyak_step_size: TensorList | Any = (state.fx0 - self.min_obj_value) / denom.where(denom!=0, 1) # type:ignore
+            polyak_step_size = polyak_step_size.where(denom != 0, 0)
+            if self.max is not None: polyak_step_size = polyak_step_size.clamp_max(self.max)
+
+        else:
+            if self.use_grad: denom = (ascent*state.maybe_compute_grad_(self.get_params())).total_mean()
+            else: denom = ascent.pow(2).total_mean()
+            if denom == 0: polyak_step_size = 0 # we converged
+            else: polyak_step_size = (state.fx0 - self.min_obj_value) / denom
+
+            if self.max is not None:
+                if polyak_step_size > self.max: polyak_step_size = self.max
+
+        ascent.mul_(alpha * polyak_step_size)
+        return ascent
+
+
+
+class RandomStepSize(OptimizerModule):
+    """Uses random global step size from `low` to `high`.
+
+    Args:
+        low (float, optional): minimum learning rate. Defaults to 0.
+        high (float, optional): maximum learning rate. Defaults to 1.
+        parameterwise (bool, optional):
+            if True, generate random step size for each parameter separately,
+            if False generate one global random step size. Defaults to False.
+    """
+    def __init__(self, low: float = 0, high: float = 1, parameterwise=False):
+        super().__init__({})
+        self.low = low; self.high = high
+        self.parameterwise = parameterwise
+
+    def _update(self, state, ascent):
+        if self.parameterwise:
+            lr = [random.uniform(self.low, self.high) for _ in range(len(ascent))]
+        else:
+            lr = random.uniform(self.low, self.high)
+        return ascent.mul_(lr) # type:ignore

torchzero/modules/second_order/__init__.py

@@ -0,0 +1,4 @@
+r"""
+This includes modules that use the hessian computed via autograd.
+"""
+from .newton import ExactNewton, LinearSystemSolvers, FallbackLinearSystemSolvers, LINEAR_SYSTEM_SOLVERS

torchzero/modules/second_order/newton.py

@@ -0,0 +1,165 @@
+from typing import Literal
+from collections import abc
+
+import torch
+
+from ...utils.derivatives import hessian_list_to_mat, jacobian_and_hessian
+from ...tensorlist import TensorList
+from ...core import OptimizerModule
+
+
+def _cholesky_solve(hessian: torch.Tensor, grad: torch.Tensor):
+    cholesky, info = torch.linalg.cholesky_ex(hessian) # pylint:disable=not-callable
+    if info == 0:
+        grad.unsqueeze_(1)
+        return torch.cholesky_solve(grad, cholesky), True
+    return None, False
+
+def _lu_solve(hessian: torch.Tensor, grad: torch.Tensor):
+    try:
+        newton_step, info = torch.linalg.solve_ex(hessian, grad) # pylint:disable=not-callable
+        if info == 0: return newton_step, True
+        return None, False
+    except torch.linalg.LinAlgError:
+        return None, False
+
+
+def _cholesky_fallback_lu(hessian: torch.Tensor, grad: torch.Tensor):
+    step, success = _cholesky_solve(hessian, grad)
+    if not success:
+        step, success = _lu_solve(hessian, grad)
+    return step, success
+
+def _least_squares_solve(hessian: torch.Tensor, grad: torch.Tensor):
+    return torch.linalg.lstsq(hessian, grad)[0], True # pylint:disable=not-callable
+
+
+def _fallback_gd(hessian:torch.Tensor, grad:torch.Tensor, lr = 1e-2):
+    return grad.mul_(1e-2), True
+
+def _fallback_safe_diag(hessian:torch.Tensor, grad:torch.Tensor, lr = 1e-2):
+    diag = hessian.diag().reciprocal_().nan_to_num_(1,1,1)
+    if torch.all(diag == 1): # fallback to gd
+        return _fallback_gd(hessian, grad, lr)
+    return grad.mul_(diag * lr), True
+
+
+def regularize_hessian_(hessian: torch.Tensor, value: float | Literal['eig']):
+    """regularize hessian matrix in-place"""
+    if value == 'eig':
+        value = torch.linalg.eigvalsh(hessian).min().clamp_(max=0).neg_() # pylint:disable=not-callable
+    elif value != 0:
+        hessian.add_(torch.eye(hessian.shape[0], device=hessian.device,dtype=hessian.dtype), alpha = value)
+
+LinearSystemSolvers = Literal['cholesky', 'lu', 'cholesky_lu', 'lstsq']
+FallbackLinearSystemSolvers = Literal['lstsq', 'safe_diag', 'gd']
+
+LINEAR_SYSTEM_SOLVERS = {
+    "cholesky": _cholesky_solve,
+    "lu": _lu_solve,
+    "cholesky_lu": _cholesky_fallback_lu,
+    "lstsq": _least_squares_solve,
+    "safe_diag": _fallback_safe_diag,
+    "gd": _fallback_gd
+}
+
+class ExactNewton(OptimizerModule):
+    """Peforms an exact Newton step using batched autograd.
+
+    Note that this doesn't support per-group settings.
+
+    Args:
+        tikhonov (float, optional):
+            tikhonov regularization (constant value added to the diagonal of the hessian).
+            Also known as Levenberg-Marquardt regularization. Can be set to 'eig', so it will be set
+            to the smallest eigenvalue of the hessian if that value is negative. Defaults to 0.
+        solver (Solvers, optional):
+            solver for Hx = g. Defaults to "cholesky_lu" (cholesky or LU if it fails).
+        fallback (Solvers, optional):
+            what to do if solver fails. Defaults to "safe_diag"
+            (takes nonzero diagonal elements, or fallbacks to gradient descent if all elements are 0).
+        validate (bool, optional):
+            validate if the step didn't increase the loss by `loss * tol` with an additional forward pass.
+            If not, undo the step and perform a gradient descent step.
+        tol (float, optional):
+            only has effect if `validate` is enabled.
+            If loss increased by `loss * tol`, perform gradient descent step.
+            Set this to 0 to guarantee that loss always decreases. Defaults to 1.
+        gd_lr (float, optional):
+            only has effect if `validate` is enabled.
+            Gradient descent step learning rate. Defaults to 1e-2.
+        batched_hessian (bool, optional):
+            whether to use experimental pytorch vmap-vectorized hessian calculation. As per pytorch docs,
+            should be faster, but this feature being experimental, there may be performance cliffs.
+            Defaults to True.
+        diag (False, optional):
+            only use the diagonal of the hessian. This will still calculate the full hessian!
+            This is mainly useful for benchmarking.
+    """
+    def __init__(
+        self,
+        tikhonov: float | Literal['eig'] = 0.0,
+        solver: LinearSystemSolvers = "cholesky_lu",
+        fallback: FallbackLinearSystemSolvers = "safe_diag",
+        validate=False,
+        tol: float = 1,
+        gd_lr = 1e-2,
+        batched_hessian=True,
+        diag: bool = False,
+    ):
+        super().__init__({})
+        self.tikhonov: float | Literal['eig'] = tikhonov
+        self.batched_hessian = batched_hessian
+
+        self.solver: abc.Callable = LINEAR_SYSTEM_SOLVERS[solver]
+        self.fallback: abc.Callable = LINEAR_SYSTEM_SOLVERS[fallback]
+
+        self.validate = validate
+        self.gd_lr = gd_lr
+        self.tol = tol
+
+        self.diag = diag
+
+    @torch.no_grad
+    def step(self, state):
+        if state.closure is None: raise ValueError("Newton requires a closure to compute the gradient.")
+
+        params = self.get_params()
+
+        # exact hessian via autograd
+        with torch.enable_grad():
+            state.fx0 = state.closure(False)
+            grads, hessian = jacobian_and_hessian([state.fx0], params) # type:ignore
+            state.grad = grads = TensorList(grads).squeeze_(0)
+            gvec = grads.to_vec()
+            hessian = hessian_list_to_mat(hessian)
+
+        # tikhonov regularization
+        regularize_hessian_(hessian, self.tikhonov)
+
+        # calculate newton step
+        if self.diag:
+            newton_step = gvec / hessian.diag()
+        else:
+            newton_step, success = self.solver(hessian, gvec)
+            if not success:
+                newton_step, success = self.fallback(hessian, gvec)
+                if not success:
+                    newton_step, success = _fallback_gd(hessian, gvec)
+
+        # apply the `_update` method
+        state.ascent = grads.from_vec(newton_step.squeeze_().nan_to_num_(0,0,0))
+
+        # validate if newton step decreased loss
+        if self.validate:
+
+            params.sub_(state.ascent)
+            fx1 = state.closure(False)
+            params.add_(state.ascent)
+
+            # if loss increases, set ascent direction to grad times lr
+            if (not fx1.isfinite()) or fx1 - state.fx0 > state.fx0 * self.tol: # type:ignore
+                state.ascent = grads.div_(grads.total_vector_norm(2) / self.gd_lr)
+
+        # peform an update with the ascent direction, or pass it to the child.
+        return self._update_params_or_step_with_next(state, params=params)

torchzero/modules/smoothing/__init__.py

@@ -0,0 +1,5 @@
+r"""
+Gradient smoothing and orthogonalization methods.
+"""
+from .laplacian_smoothing import LaplacianSmoothing, gradient_laplacian_smoothing_
+from .gaussian_smoothing import GaussianSmoothing

torchzero/modules/smoothing/gaussian_smoothing.py

@@ -0,0 +1,90 @@
+from contextlib import nullcontext
+
+import numpy as np
+import torch
+
+from ...tensorlist import TensorList, Distributions, mean as tlmean
+from ...utils.python_tools import _ScalarLoss
+from ...core import _ClosureType, OptimizationState, OptimizerModule, _maybe_pass_backward
+
+
+def _numpy_or_torch_mean(losses: list):
+    """Returns the mean of a list of losses, which can be either numpy arrays or torch tensors."""
+    if isinstance(losses[0], torch.Tensor):
+        return torch.mean(torch.stack(losses))
+    return np.mean(losses).item()
+
+class GaussianSmoothing(OptimizerModule):
+    """Samples and averages value and gradients in multiple random points around current position.
+    This effectively applies smoothing to the function.
+
+    Args:
+        n_samples (int, optional): number of gradient samples from around current position. Defaults to 4.
+        sigma (float, optional): how far from current position to sample from. Defaults to 0.1.
+        distribution (tl.Distributions, optional): distribution for random positions. Defaults to "normal".
+        sample_x0 (bool, optional): 1st sample will be x0. Defaults to False.
+        randomize_every (int | None, optional): randomizes the points every n steps. Defaults to 1.
+    """
+    def __init__(
+        self,
+        n_samples: int = 4,
+        sigma: float = 0.1,
+        distribution: Distributions = "normal",
+        sample_x0 = False,
+        randomize_every: int | None = 1,
+    ):
+        defaults = dict(sigma = sigma)
+        super().__init__(defaults)
+        self.n_samples = n_samples
+        self.distribution: Distributions = distribution
+        self.randomize_every = randomize_every
+        self.current_step = 0
+        self.perturbations = None
+        self.sample_x0 = sample_x0
+
+    @torch.no_grad()
+    def step(self, state: OptimizationState):
+        if state.closure is None: raise ValueError('GaussianSmoothing requires closure.')
+        closure = state.closure
+        params = self.get_params()
+        sigmas = self.get_group_key('sigma')
+
+        # generate random perturbations
+        if self.perturbations is None or (self.randomize_every is not None and self.current_step % self.randomize_every == 0):
+            if self.sample_x0:
+                self.perturbations = [params.sample_like(sigmas, distribution=self.distribution) for _ in range(self.n_samples-1)]
+            else:
+                self.perturbations = [params.sample_like(sigmas, distribution=self.distribution) for _ in range(self.n_samples)]
+
+        @torch.no_grad
+        def smooth_closure(backward = True):
+            losses = []
+            grads = []
+
+            # sample gradient and loss at x0
+            if self.sample_x0:
+                with torch.enable_grad() if backward else nullcontext():
+                    losses.append(closure())
+                    if backward: grads.append(params.grad.clone())
+
+            # sample gradients from points around current params
+            # and average them
+            if self.perturbations is None: raise ValueError('who set perturbations to None???')
+            for p in self.perturbations:
+                params.add_(p)
+                with torch.enable_grad() if backward else nullcontext():
+                    losses.append(_maybe_pass_backward(closure, backward))
+                    if backward: grads.append(params.grad.clone())
+                params.sub_(p)
+
+            # set the new averaged grads and return average loss
+            if backward: params.set_grad_(tlmean(grads))
+            return _numpy_or_torch_mean(losses)
+
+
+        self.current_step += 1
+        state.closure = smooth_closure
+        return self._update_params_or_step_with_next(state)
+
+
+# todo single loop gaussian homotopy?

torchzero/modules/smoothing/laplacian_smoothing.py

@@ -0,0 +1,128 @@
+from typing import Literal
+from collections.abc import Iterable
+
+import torch
+
+from ...tensorlist import TensorList
+from ...core import OptimizerModule
+
+
+def vector_laplacian_smoothing(input: torch.Tensor, sigma: float = 1) -> torch.Tensor:
+    """Returns a new vector with laplacian smoothing applied to it. This flattens the input!"""
+    vec = input.view(-1)
+    v = torch.zeros_like(vec)
+    v[0] = -2
+    v[1] = 1
+    v[-1] = 1
+    numerator = torch.fft.fft(vec) # pylint: disable = not-callable
+    denominator = 1 - sigma * torch.fft.fft(v) # pylint: disable = not-callable
+    return torch.fft.ifft(numerator / denominator).real # pylint: disable = not-callable
+
+def gradient_laplacian_smoothing_(params: Iterable[torch.Tensor], sigma: float = 1, layerwise=True, min_numel = 4):
+    """Applies laplacian smoothing to gradients of an iterable of parameters.
+
+    This updates gradients in-place.
+
+    Args:
+        params (abc.Iterable[torch.Tensor]): an iterable of Tensors that will have gradients smoothed.
+        sigma (float, optional): controls the amount of smoothing. Defaults to 1.
+        layerwise (bool, optional):
+            If True, applies smoothing to each parameter's gradient separately,
+            Otherwise applies it to all gradients, concatenated into a single vector. Defaults to True.
+        min_numel (int, optional):
+            minimum number of elements in a parameter to apply laplacian smoothing to.
+            Only has effect if `layerwise` is True. Defaults to 4.
+
+    Reference:
+        *Osher, S., Wang, B., Yin, P., Luo, X., Barekat, F., Pham, M., & Lin, A. (2022).
+        Laplacian smoothing gradient descent. Research in the Mathematical Sciences, 9(3), 55.*
+    """
+    grads = TensorList(params).get_existing_grads()
+    if layerwise:
+        for g in grads:
+            if g.numel() >= min_numel:
+                g.set_(vector_laplacian_smoothing(g, sigma).reshape(g.shape)) # type:ignore
+    else:
+        vec = grads.to_vec()
+        grads.from_vec_(vector_laplacian_smoothing(vec, sigma))
+
+
+def _precompute_denominator(tensor: torch.Tensor, sigma) -> torch.Tensor:
+    """Denominator will always be the same and depends on the size of the vector and the sigma."""
+    v = torch.zeros_like(tensor.view(-1))
+    v[0] = -2
+    v[1] = 1
+    v[-1] = 1
+    return 1 - sigma * torch.fft.fft(v) # pylint: disable = not-callable
+
+class LaplacianSmoothing(OptimizerModule):
+    """Applies laplacian smoothing via a fast Fourier transform solver.
+
+    Args:
+        sigma (float, optional): controls the amount of smoothing. Defaults to 1.
+        layerwise (bool, optional):
+            If True, applies smoothing to each parameter's gradient separately,
+            Otherwise applies it to all gradients, concatenated into a single vector. Defaults to True.
+        min_numel (int, optional):
+            minimum number of elements in a parameter to apply laplacian smoothing to.
+            Only has effect if `layerwise` is True. Defaults to 4.
+        target (str, optional):
+            determines what this module updates.
+
+            "ascent" - it updates the ascent (default).
+
+            "grad" - it updates the gradient (and sets `.grad` attributes to updated gradient).
+
+            "closure" - it makes a new closure that sets the updated ascent to the .`grad` attributes.
+
+    Reference:
+        *Osher, S., Wang, B., Yin, P., Luo, X., Barekat, F., Pham, M., & Lin, A. (2022).
+        Laplacian smoothing gradient descent. Research in the Mathematical Sciences, 9(3), 55.*
+
+    """
+    def __init__(self, sigma:float = 1, layerwise=True, min_numel = 4, target: Literal['ascent', 'grad', 'closure',] = 'ascent'):
+        # sigma from defaults is used in layerwise case
+        # otherwise self.sigma is used
+        defaults = dict(sigma = sigma)
+        self.sigma = 1
+        super().__init__(defaults, target=target)
+        self.layerwise = layerwise
+        self.min_numel = min_numel
+
+        # precomputed denominator for when layerwise=False
+        self.full_denominator = None
+
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        params = self.get_params()
+        sigmas = self.get_group_key('sigma')
+
+        # layerwise laplacian smoothing
+        if self.layerwise:
+
+            # precompute the denominator for each layer and store it in each parameters state
+            denominators = TensorList()
+            for p, σ in zip(params, sigmas):
+                if p.numel() > self.min_numel:
+                    den = self.state[p]
+                    if 'denominator' not in den: den['denominator'] = _precompute_denominator(p, σ)
+                    denominators.append(den['denominator'])
+
+            # apply the smoothing
+            smoothed_direction = TensorList()
+            for g, σ, den in zip(ascent, sigmas, denominators):
+                smoothed_direction.append(torch.fft.ifft(torch.fft.fft(g.view(-1)) / den).real.reshape(g.shape)) # pylint: disable = not-callable
+            return smoothed_direction
+
+        # else
+        # full laplacian smoothing
+        # precompute full denominator
+        if self.full_denominator is None:
+            self.full_denominator = _precompute_denominator(ascent.to_vec(), self.sigma)
+
+        # apply the smoothing
+        vec = ascent.to_vec()
+        return ascent.from_vec(torch.fft.ifft(torch.fft.fft(vec) / self.full_denominator).real) # pylint: disable = not-callable
+
+

torchzero/modules/weight_averaging/__init__.py

@@ -0,0 +1,2 @@
+from .ema import SwitchEMA
+from .swa import PeriodicSWA, CyclicSWA

torchzero/modules/weight_averaging/ema.py

@@ -0,0 +1,72 @@
+import torch
+from ...core import OptimizerModule
+
+
+def _reset_stats_hook(optimizer, state):
+    for module in optimizer.unrolled_modules:
+        module: OptimizerModule
+        module.reset_stats()
+
+# the reason why this needs to be at the end is ??? I NEED TO REMEMBER
+class SwitchEMA(OptimizerModule):
+    """Switch-EMA. Every n steps switches params to an exponential moving average of past weights.
+
+    In the paper the switch happens after each epoch.
+
+    Please put this module at the end, after all other modules.
+
+    This can also function as EMA, set `update_every` to None and instead call `set_ema` and `unset_ema` on this module.
+
+
+    Args:
+        update_every (int): number of steps (batches) between setting model parameters to EMA.
+        momentum (int): EMA momentum factor.
+        reset_stats (bool, optional):
+            if True, when setting model parameters to EMA, resets other modules stats such as momentum velocities.
+            It might be better to set this to False if `update_every` is very small. Defaults to True.
+
+    reference
+        https://arxiv.org/abs/2402.09240
+    """
+    def __init__(self, update_every: int | None, momentum: float = 0.99, reset_stats: bool = True):
+        defaults = dict(momentum=momentum)
+        super().__init__(defaults)
+        self.update_every = update_every
+        self.cur_step = 0
+        self.update_every = update_every
+        self._reset_stats = reset_stats
+        self.orig_params = None
+
+    def set_ema(self):
+        """sets module parameters to EMA, stores original parameters that can be restored by calling `unset_ema`"""
+        params = self.get_params()
+        self.orig_params = params.clone()
+        params.set_(self.get_state_key('ema', init = 'params', params=params))
+
+    def unset_ema(self):
+        """Undoes `set_ema`."""
+        if self.orig_params is None: raise ValueError('call `set_ema` first, and then `unset_ema`.')
+        params = self.get_params()
+        params.set_(self.orig_params)
+
+    @torch.no_grad
+    def step(self, state):
+        # if self.next_module is not None:
+        #     warn(f'EMA should usually be the last module, but {self.next_module.__class__.__name__} is after it.')
+        self.cur_step += 1
+
+        params = self.get_params()
+        # state.maybe_use_grad_(params)
+        # update params with the child. Averaging is always applied at the end.
+        ret = self._update_params_or_step_with_next(state, params)
+
+        ema = self.get_state_key('ema', init = 'params', params=params)
+        momentum = self.get_group_key('momentum')
+
+        ema.lerp_compat_(params, 1 - momentum)
+
+        if (self.update_every is not None) and (self.cur_step % self.update_every == 0):
+            params.set_(ema.clone())
+            if self._reset_stats: state.add_post_step_hook(_reset_stats_hook)
+
+        return ret