torchzero 0.1.8__py3-none-any.whl → 0.3.2__py3-none-any.whl

torchzero/modules/operations/reduction.py

@@ -1,134 +0,0 @@
-from collections.abc import Callable, Iterable
-import numpy as np
-import torch
-
-from ...core import OptimizerModule
-
-_Value = int | float | OptimizerModule | Iterable[OptimizerModule]
-
-
-class Sum(OptimizerModule):
-    """calculates sum of multiple updates.
-
-    Args:
-        *modules:
-            either OptimizerModules or iterables of OptimizerModules to chain. Scalars are also allowed."""
-    def __init__(
-        self,
-        *modules: _Value,
-    ):
-        super().__init__({})
-
-        scalars = [i for i in modules if isinstance(i, (int,float))]
-        self.scalar = sum(scalars) if len(scalars) > 0 else None
-
-        for i,module in enumerate(i for i in modules if not isinstance(i, (int, float))):
-            self._set_child_(i, module)
-
-    @torch.no_grad
-    def step(self, vars):
-        if len(self.children) == 1:
-            vars.ascent = self.children[0].return_ascent(vars)
-            if self.scalar is not None: vars.ascent += self.scalar
-            return self._update_params_or_step_with_next(vars)
-
-        sum = None
-        for i, c in sorted(self.children.items(), key=lambda x: x[0]):
-            if i == len(self.children) - 1: cur_state = vars
-            else: cur_state = vars.copy(clone_ascent = True)
-
-            if sum is None: sum = c.return_ascent(cur_state)
-            else: sum += c.return_ascent(cur_state)
-
-            if i != len(self.children) - 1: vars.update_attrs_(cur_state)
-
-        assert sum is not None
-        if self.scalar is not None: sum += self.scalar
-        vars.ascent = sum
-        return self._update_params_or_step_with_next(vars)
-
-class Mean(OptimizerModule):
-    """calculates mean of multiple updates.
-
-    Args:
-        *modules:
-            either OptimizerModules or iterables of OptimizerModules to chain. Scalars are also allowed."""
-
-    def __init__(
-        self,
-        *modules: _Value,
-    ):
-        super().__init__({})
-
-        scalars = [i for i in modules if isinstance(i, (int,float))]
-        self.scalar = sum(scalars) if len(scalars) > 0 else None
-
-        self.n_values = len(modules)
-
-        for i,module in enumerate(i for i in modules if not isinstance(i, (int, float))):
-            self._set_child_(i, module)
-
-    @torch.no_grad
-    def step(self, vars):
-        if len(self.children) == 1:
-            vars.ascent = self.children[0].return_ascent(vars)
-            if self.scalar is not None: vars.ascent += self.scalar
-            if self.n_values > 1: vars.ascent /= self.n_values
-            return self._update_params_or_step_with_next(vars)
-
-        sum = None
-        for i, c in sorted(self.children.items(), key=lambda x: x[0]):
-            if i == len(self.children) - 1: cur_state = vars
-            else: cur_state = vars.copy(clone_ascent = True)
-
-            if sum is None: sum = c.return_ascent(cur_state)
-            else: sum += c.return_ascent(cur_state)
-
-            if i != len(self.children) - 1: vars.update_attrs_(cur_state)
-
-        assert sum is not None
-        if self.scalar is not None: sum += self.scalar
-        if self.n_values > 1: sum /= self.n_values
-        vars.ascent = sum
-        return self._update_params_or_step_with_next(vars)
-
-class Product(OptimizerModule):
-    """calculates product of multiple updates.
-
-    Args:
-        *modules:
-            either OptimizerModules or iterables of OptimizerModules to chain. Scalars are also allowed."""
-
-    def __init__(
-        self,
-        *modules: _Value,
-    ):
-        super().__init__({})
-
-        scalars = [i for i in modules if isinstance(i, (int,float))]
-        self.scalar = np.prod(scalars).item() if len(scalars) > 0 else None
-
-        for i,module in enumerate(i for i in modules if not isinstance(i, (int, float))):
-            self._set_child_(i, module)
-
-    @torch.no_grad
-    def step(self, vars):
-        if len(self.children) == 1:
-            vars.ascent = self.children[0].return_ascent(vars)
-            if self.scalar is not None: vars.ascent *= self.scalar
-            return self._update_params_or_step_with_next(vars)
-
-        prod = None
-        for i, c in sorted(self.children.items(), key=lambda x: x[0]):
-            if i == len(self.children) - 1: cur_state = vars
-            else: cur_state = vars.copy(clone_ascent = True)
-
-            if prod is None: prod = c.return_ascent(cur_state)
-            else: prod *= c.return_ascent(cur_state)
-
-            if i != len(self.children) - 1: vars.update_attrs_(cur_state)
-
-        assert prod is not None
-        if self.scalar is not None: prod *= self.scalar
-        vars.ascent = prod
-        return self._update_params_or_step_with_next(vars)

torchzero/modules/operations/singular.py

@@ -1,113 +0,0 @@
-from collections.abc import Iterable
-from operator import methodcaller
-
-import torch
-
-from ...core import OptimizerModule
-from ...tensorlist import TensorList
-
-
-class Operation(OptimizerModule):
-    """Applies an operation to the ascent, supported operations:
-
-    `abs`, `sign`, `sin`, `cos`, `tan`, `asin`, `acos`, `atan`, `sinh`, `cosh`,
-    `tanh`, `log`, `log1p`, `log2`, `log10`, `erf`, `erfc`, `exp`, `neg`, `reciprocal`,
-    `copy`, `zero`, `sqrt`, `floor`, `ceil`, `round`."""
-    def __init__(self, operation: str):
-        super().__init__({})
-        self.operation = methodcaller(f'{operation}_')
-
-    @torch.no_grad
-    def _update(self, vars, ascent): return self.operation(ascent)
-
-class Reciprocal(OptimizerModule):
-    """*1 / update*"""
-    def __init__(self,):
-        super().__init__({})
-
-    @torch.no_grad()
-    def _update(self, vars, ascent): return ascent.reciprocal_()
-
-class Negate(OptimizerModule):
-    """minus update"""
-    def __init__(self,):
-        super().__init__({})
-
-    @torch.no_grad()
-    def _update(self, vars, ascent): return ascent.neg_()
-
-
-def sign_grad_(params: Iterable[torch.Tensor]):
-    """Apply sign function to gradients of an iterable of parameters.
-
-    Args:
-        params (abc.Iterable[torch.Tensor]): an iterable of Tensors or a single Tensor.
-    """
-    TensorList(params).get_existing_grads().sign_()
-
-class Sign(OptimizerModule):
-    """applies sign function to the update"""
-    def __init__(self):
-        super().__init__({})
-
-    @torch.no_grad
-    def _update(self, vars, ascent): return ascent.sign_()
-
-class Abs(OptimizerModule):
-    """takes absolute values of the update."""
-    def __init__(self):
-        super().__init__({})
-
-    @torch.no_grad
-    def _update(self, vars, ascent): return ascent.abs_()
-
-class Sin(OptimizerModule):
-    """applies sin function to the ascent"""
-    def __init__(self):
-        super().__init__({})
-
-    @torch.no_grad
-    def _update(self, vars, ascent): return ascent.sin_()
-
-class Cos(OptimizerModule):
-    """applies cos function to the ascent"""
-    def __init__(self):
-        super().__init__({})
-
-    @torch.no_grad
-    def _update(self, vars, ascent): return ascent.cos_()
-
-
-class NanToNum(OptimizerModule):
-    """Convert `nan`, `inf` and `-inf` to numbers.
-
-    Args:
-        nan (optional): the value to replace NaNs with. Default is zero.
-        posinf (optional): if a Number, the value to replace positive infinity values with.
-            If None, positive infinity values are replaced with the greatest finite value
-            representable by input's dtype. Default is None.
-        neginf (optional): if a Number, the value to replace negative infinity values with.
-            If None, negative infinity values are replaced with the lowest finite value
-            representable by input's dtype. Default is None.
-    """
-    def __init__(self, nan=None, posinf=None, neginf=None):
-        super().__init__({})
-        self.nan = nan
-        self.posinf = posinf
-        self.neginf = neginf
-
-    @torch.no_grad()
-    def _update(self, vars, ascent): return ascent.nan_to_num_(self.nan, self.posinf, self.neginf)
-
-
-class MagnitudePower(OptimizerModule):
-    """Raises update to the `value` power, but preserves the sign when the power is odd."""
-    def __init__(self, value: int | float):
-        super().__init__({})
-        self.value = value
-
-    @torch.no_grad()
-    def _update(self, vars, ascent):
-        if self.value % 2 == 1: return ascent.pow_(self.value)
-        return ascent.abs().pow_(self.value) * ascent.sign()
-

torchzero/modules/optimizers/sgd.py

@@ -1,54 +0,0 @@
-import typing as T
-
-import torch
-
-from ...core import OptimizerModule
-from ..momentum.momentum import _heavyball_step, _nesterov_step_
-
-class SGD(OptimizerModule):
-    """Same as `torch.optim.SGD` but as an optimizer module. Exactly matches `torch.optim.SGD`, except
-    nesterov momentum additionally supports dampening, and negative momentum is allowed.
-
-    Args:
-        momentum (float, optional): momentum. Defaults to 0.
-        dampening (float, optional): momentum dampening. Defaults to 0.
-        weight_decay (float, optional): weight decay (L2 regularization). Defaults to 0.
-        nesterov (bool, optional):
-            enables nesterov momentum, otherwise uses heavyball momentum. Defaults to False.
-        alpha (float, optional): learning rate. Defaults to 1.
-    """
-    def __init__(
-        self,
-        momentum: float = 0,
-        dampening: float = 0,
-        weight_decay: float = 0,
-        nesterov: bool = False,
-        alpha: float = 1,
-    ):
-
-        defaults = dict(alpha=alpha, momentum=momentum, dampening=dampening, weight_decay=weight_decay,)
-        super().__init__(defaults)
-        self.nesterov = nesterov
-        self.current_step = 0
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        params = self.get_params()
-        settings = self.get_all_group_keys()
-
-        if any(i != 0 for i in settings['weight_decay']):
-            ascent += params * settings['weight_decay']
-
-        if any(i != 1 for i in settings['alpha']):
-            ascent *= settings['alpha']
-
-        if any(i != 0 for i in settings['momentum']):
-            velocity = self.get_state_key('velocity', init = torch.zeros_like if self.nesterov else ascent)
-            # consistency with pytorch which on first step only initializes momentum
-            if self.current_step > 0 or self.nesterov:
-                # nesterov step can be done in-place, polyak returns new direction
-                if self.nesterov: _nesterov_step_(ascent, velocity, settings['momentum'], settings['dampening'])
-                else: ascent = _heavyball_step(ascent, velocity, settings['momentum'], settings['dampening'])
-
-        self.current_step += 1
-        return ascent

torchzero/modules/orthogonalization/__init__.py

@@ -1,2 +0,0 @@
-from .svd import Orthogonalize, orthogonalize_grad_
-from .newtonschulz import ZeropowerViaNewtonSchulz, zeropower_via_newtonschulz_, DualNormCorrection

torchzero/modules/orthogonalization/newtonschulz.py

@@ -1,159 +0,0 @@
-"""
-Newton-Schulz iteration code is taken from https://github.com/KellerJordan/Muon
-
-Keller Jordan and Yuchen Jin and Vlado Boza and You Jiacheng and Franz Cecista and Laker Newhouse and Jeremy Bernstein.
-Muon: An optimizer for hidden layers in neural networks (2024). URL: https://kellerjordan.github.io/posts/muon
-"""
-from collections.abc import Iterable
-
-import torch
-
-from ...core import OptimizerModule, _Targets
-# from ...utils.compile import maybe_compile
-
-def _zeropower_via_newtonschulz5(G, steps):
-    """
-    code from https://github.com/KellerJordan/Muon
-
-    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
-    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
-    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
-    zero even beyond the point where the iteration no longer converges all the way to one everywhere
-    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
-    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
-    performance at all relative to UV^T, where USV^T = G is the SVD.
-    """
-    assert len(G.shape) == 2
-    a, b, c = (3.4445, -4.7750,  2.0315)
-    X = G.bfloat16()
-    if G.size(0) > G.size(1):
-        X = X.T
-
-    # Ensure spectral norm is at most 1
-    X = X / (X.norm() + 1e-7)
-    # Perform the NS iterations
-    for _ in range(steps):
-        A = X @ X.T
-        B = b * A + c * A @ A # adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
-        X = a * X + B @ X
-
-    if G.size(0) > G.size(1):
-        X = X.T
-
-    return X
-
-_compiled_zeropower_via_newtonschulz5 = torch.compile(_zeropower_via_newtonschulz5)
-
-
-def zeropower_via_newtonschulz_(params: Iterable[torch.Tensor], steps: int = 6, adaptive = False, compiled = True):
-    """Uses newton-Schulz iteration to compute the zeroth power / orthogonalization of gradients of an iterable of parameters.
-
-    This sets gradients in-place.
-
-    Note that the Muon page says that embeddings and classifier heads should not be orthogonalized.
-
-    The orthogonalization code is taken from https://github.com/KellerJordan/Muon
-    Args:
-        params (abc.Iterable[torch.Tensor]): parameters that hold gradients to orthogonalize.
-        steps (int): The number of Newton-Schulz iterations to run. (6 is probably always enough).
-            The number of Newton-Schulz iterations to run. (6 is probably always enough). Defaults to 6.
-        adaptive (bool, optional):
-            Enables adaptation to scale of gradients (from https://github.com/leloykun/adaptive-muon). Defaults to False.
-        compiled (bool, optional):
-            Uses compiled newton-Schulz iteration function. Faster but won't work on windows. Defaults to True.
-
-
-    """
-    if compiled: fn = _compiled_zeropower_via_newtonschulz5
-    else: fn = _zeropower_via_newtonschulz5
-    for p in params:
-        if p.grad is not None and p.grad.ndim >= 2 and min(p.grad.shape) >= 2:
-            G = p.grad.view(p.grad.shape[0], -1)
-            X = fn(G, steps)
-
-            if adaptive:
-                # this is from https://github.com/leloykun/adaptive-muon
-                X = torch.einsum('ij,ij,ab->ab', G.type_as(X), X, X)  # Adaptive scaling,`(G * X).sum() * X` == (G.T @ X).trace() * X
-
-            p.grad = X.reshape_as(p.grad).to(p.grad, copy=False)
-
-
-class ZeropowerViaNewtonSchulz(OptimizerModule):
-    """Uses Newton-Schulz iteration to compute the zeroth power / orthogonalization of gradients of an iterable of parameters.
-
-    To disable orthogonalization for a parameter, put it into a parameter group with "newtonshultz" = False.
-    The Muon page says that embeddings and classifier heads should not be orthogonalized.
-
-    The orthogonalization code is taken from https://github.com/KellerJordan/Muon.
-
-    Note that unlike this module, Muon also uses Adam for gradients that are not orthogonalized,
-    so I'd still recommend using it. Maybe use `Wrap` to wrap it into a module (I will make muon
-    with selectable modules to optimize non-muon params soon)
-
-    However not using Adam, or putting Adam module after this to apply it to ALL updates, both seem
-    to work quite well too.
-
-    Args:
-        ns_steps (int, optional):
-            The number of Newton-Schulz iterations to run. (6 is probably always enough). Defaults to 6.
-        adaptive (bool, optional):
-            Enables adaptation to scale of gradients (from https://github.com/leloykun/adaptive-muon). Defaults to True.
-        compiled (bool, optional):
-            Uses compiled newton-Schulz iteration function. Faster but won't work on windows. Defaults to True.
-        target (str, optional):
-            determines what this module updates.
-
-            "ascent" - it updates the ascent
-
-            "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.
-    """
-    def __init__(self, ns_steps = 6, adaptive = False, compiled=True, target:_Targets='ascent'):
-        defaults = dict(newtonshultz = True, ns_steps=ns_steps, adaptive=adaptive)
-        super().__init__(defaults, target=target)
-
-        if compiled: self._zeropower_via_newtonschulz5 = _compiled_zeropower_via_newtonschulz5
-        else: self._zeropower_via_newtonschulz5 = _zeropower_via_newtonschulz5
-
-    def _update(self, vars, ascent):
-        toggle, ns_steps, adaptive = self.get_group_keys('newtonshultz', 'ns_steps', 'adaptive', cls=list)
-
-        for asc, enable, steps, ada in zip(ascent, toggle, ns_steps, adaptive):
-            if enable and len([i for i in asc.shape if i > 1]) != 0:
-                G = asc.view(asc.shape[0], -1)
-                X = self._zeropower_via_newtonschulz5(G, steps)
-
-                if ada:
-                    # this is from https://github.com/leloykun/adaptive-muon
-                    X = torch.einsum('ij,ij,ab->ab', G.type_as(X), X, X)  # Adaptive scaling,`(G * X).sum() * X` == (G.T @ X).trace() * X
-
-                asc.set_(X.reshape_as(asc).to(asc, copy=False)) # type:ignore
-
-        return ascent
-
-
-
-class DualNormCorrection(OptimizerModule):
-    """Dual norm correction from https://github.com/leloykun/adaptive-muon.
-
-    Description from the page:
-
-    Single-line modification to any (dualizer-based) optimizer that allows the optimizer to adapt to the scale of the gradients as they change during training.
-    This is done by scaling the dualized gradient by the clipped dual norm of the original gradient.
-    """
-    def __init__(self, adaptive_scale_min: int | None = -1, adaptive_scale_max: int | None = 1):
-        defaults = dict(adaptive_scale_min = adaptive_scale_min, adaptive_scale_max = adaptive_scale_max)
-        super().__init__(defaults)
-
-    def _update(self, vars, ascent):
-        params = self.get_params()
-        adaptive_scale_min, adaptive_scale_max = self.get_group_keys('adaptive_scale_min', 'adaptive_scale_max')
-
-        for asc, grad, min, max in zip(ascent, vars.maybe_compute_grad_(params), adaptive_scale_min, adaptive_scale_max):
-            if len([i for i in asc.shape if i > 1]) != 0:
-                scale = torch.einsum('ij,ij->', grad.view(grad.shape[0], -1), asc.view(asc.shape[0], -1))
-                if min is not None or max is not None: scale = scale.clip(min, max)
-                asc *= scale
-
-        return ascent

torchzero/modules/orthogonalization/svd.py

@@ -1,86 +0,0 @@
-"""Orthogonalization code adapted from https://github.com/MarkTuddenham/Orthogonal-Optimisers
-
-Tuddenham, M., Prügel-Bennett, A., & Hare, J. (2022).
-Orthogonalising gradients to speed up neural network optimisation. arXiv preprint arXiv:2202.07052.
-"""
-import logging
-from collections.abc import Iterable, Sequence
-
-import torch
-
-from ...core import OptimizerModule, _Targets
-
-@torch.no_grad()
-def _orthogonalize_update_(updates: Sequence[torch.Tensor], toggle = None, warn_fail=True) -> None:
-    """adapted from https://github.com/MarkTuddenham/Orthogonal-Optimisers"""
-    if toggle is None: toggle = [True] * len(updates)
-
-    # Orthogonalise the gradients using SVD
-    for grad, orth in zip(updates, toggle):
-        if orth and grad.ndim > 1:
-            G: torch.Tensor = grad.view(grad.shape[0], -1)
-            orth_G: torch.Tensor | None = None
-            try:
-                u, s, vt = torch.linalg.svd(G, full_matrices=False) # pylint:disable=not-callable
-                orth_G = u @ vt
-            except RuntimeError:
-                # if warn: logging.warning('Failed to perform SVD, adding some noise.')
-                try:
-                    u, s, v = torch.svd_lowrank(
-                        G,
-                        q=1,    # assume rank is at least 1
-                        M=1e-4 * G.mean() * torch.randn_like(G))
-                    orth_G = u @ v.T
-                except RuntimeError:
-                    if warn_fail: logging.error(('Failed to perform SVD with noise,'
-                                    ' skipping gradient orthogonalisation'))
-            if orth_G is not None:
-                grad.set_(orth_G.reshape_as(grad)) # type:ignore
-
-    return updates
-
-def orthogonalize_grad_(params: Iterable[torch.Tensor], warn_fail=False):
-    """orthogonalizes gradients of an iterable of parameters.
-
-    This updates gradients in-place.
-
-    The orthogonalization code is adapted from https://github.com/MarkTuddenham/Orthogonal-Optimisers
-    Args:
-        params (abc.Iterable[torch.Tensor]): parameters that hold gradients to orthogonalize.
-        warn_fail (bool, optional):
-            whether to print a warning when orthogonalization fails, and gradients are not
-            orthogonalized. Defaults to True.
-    """
-    grads = [p.grad for p in params if p.grad is not None]
-    _orthogonalize_update_(grads, warn_fail=warn_fail)
-
-class Orthogonalize(OptimizerModule):
-    """Orthogonalizes the update using SVD.
-
-    To disable orthogonalization for a parameter, put it into a parameter group with "orth" = False.
-
-    The orthogonalization code is adapted from https://github.com/MarkTuddenham/Orthogonal-Optimisers
-
-    Tip: :py:class:`tz.m.ZeropowerViaNewtonSchulz` is a significantly faster version of this.
-    Args:
-        warn_fail (bool, optional):
-            whether to print a warning when orthogonalization fails, and gradients are not
-            orthogonalized. Defaults to True.
-        target (str, optional):
-            determines what this module updates.
-
-            "ascent" - it updates the ascent
-
-            "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.
-    """
-    def __init__(self, warn_fail=True, target: _Targets = 'ascent'):
-        defaults = dict(orth = True)
-        super().__init__(defaults, target = target)
-        self.warn_fail = warn_fail
-
-    def _update(self, vars, ascent):
-        toggle = self.get_group_key('orth', cls=list)
-        _orthogonalize_update_(ascent, toggle, self.warn_fail)
-        return ascent

torchzero/modules/regularization/__init__.py

@@ -1,22 +0,0 @@
-r"""
-This includes regularization modules like weight decay.
-"""
-from .dropout import Dropout
-from .noise import AddNoise, Random, add_noise_
-from .normalization import (
-    Centralize,
-    ClipNorm,
-    ClipValue,
-    Normalize,
-    centralize_grad_,
-    clip_grad_norm_,
-    clip_grad_value_,
-    normalize_grad_,
-)
-from .weight_decay import (
-    WeightDecay,
-    l1_regularize_,
-    l2_regularize_,
-    weight_decay_penalty,
-)
-from .ortho_grad import OrthoGrad, orthograd_

torchzero/modules/regularization/dropout.py

@@ -1,34 +0,0 @@
-import typing as T
-from collections import abc
-
-import torch
-
-from ...tensorlist import Distributions, TensorList
-from ...core import OptimizerModule
-
-
-class Dropout(OptimizerModule):
-    """
-    Applies dropout to the update - sets random elements to 0.
-
-    This can be used to apply learning rate dropout, if put after other modules, or gradient dropout,
-    if put first.
-
-    Args:
-        p (float, optional): probability to replace update value with zero. Defaults to 0.5.
-
-    reference
-        *Lin, H., Zeng, W., Zhuang, Y., Ding, X., Huang, Y., & Paisley, J. (2022).
-        Learning rate dropout. IEEE Transactions on Neural Networks and Learning Systems,
-        34(11), 9029-9039.*
-    """
-    def __init__(self, p: float = 0.5):
-        defaults = dict(p = p)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        p = self.get_group_key('p')
-
-        ascent *= ascent.bernoulli_like(p)
-        return ascent

torchzero/modules/regularization/noise.py

@@ -1,77 +0,0 @@
-from collections import abc
-from typing import Literal
-
-import torch
-
-from ...core import OptimizerModule
-from ...tensorlist import Distributions, TensorList, _Scalar, _ScalarSequence
-
-
-def add_noise_(
-    grads: abc.Iterable[torch.Tensor],
-    alpha: "_Scalar | _ScalarSequence" = 1e-2,
-    distribution: Distributions = "normal",
-    mode: Literal["absolute", "global", "param", "channel"] = "param",
-):
-    if not isinstance(grads, TensorList): grads = TensorList(grads)
-    if mode == 'absolute':
-        grads += grads.sample_like(alpha, distribution)
-
-    elif mode == 'global':
-        grads += grads.sample_like((grads.total_vector_norm(1)/grads.total_numel() * alpha).detach().cpu().item(), distribution) # type:ignore
-
-    elif mode == 'param':
-        grads += grads.sample_like(grads.abs().mean()*alpha, distribution)
-
-    elif mode == 'channel':
-        grads = grads.unbind_channels()
-        grads += grads.sample_like(grads.abs().mean()*alpha, distribution)
-
-class AddNoise(OptimizerModule):
-    """Add noise to update. By default noise magnitude is relative to the mean of each parameter.
-
-    Args:
-        alpha (float, optional): magnitude of noise. Defaults to 1e-2.
-        distribution (Distributions, optional): distribution of noise. Defaults to 'normal'.
-        mode (str, optional):
-            how to calculate noise magnitude.
-
-            - "absolute": ignores gradient magnitude and always uses `alpha` as magnitude.
-
-            - "global": multiplies `alpha` by mean of the entire gradient, as if it was a single vector.
-
-            - "param": multiplies `alpha` by mean of each individual parameter (default).
-
-            - "channel": multiplies `alpha` by mean of each channel of each parameter.
-        """
-
-    def __init__(
-        self,
-        alpha: float = 1.,
-        distribution: Distributions = "normal",
-        mode: Literal["absolute", "global", "param", "channel"] = "param",
-    ):
-        defaults = dict(alpha = alpha)
-        super().__init__(defaults)
-        self.distribution: Distributions = distribution
-        self.mode: Literal["absolute", "global", "param", "channel"] = mode
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        alpha = self.get_group_key('alpha')
-
-        add_noise_(ascent, alpha, self.distribution, self.mode)
-        return ascent
-
-class Random(OptimizerModule):
-    """uses a random vector as the update. The vector is completely random and isn't checked to be descent direction.
-    This is therefore mainly useful in combination with other modules like Sum, Multiply, etc."""
-    def __init__(self, alpha: float = 1, distribution: Distributions = "normal"):
-        defaults = dict(alpha = alpha)
-        super().__init__(defaults)
-        self.distribution: Distributions = distribution
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        alpha = self.get_group_key('alpha')
-        return ascent.sample_like(alpha, self.distribution)