torchzero 0.3.8__py3-none-any.whl → 0.3.10__py3-none-any.whl

torchzero/modules/line_search/line_search.py

@@ -8,7 +8,7 @@ from typing import Any
 import numpy as np
 import torch
 
-from ...core import Module, Target, Vars
+from ...core import Module, Target, Var
 from ...utils import tofloat
 
 
@@ -62,12 +62,12 @@ class LineSearch(Module, ABC):
         if any(a!=0 for a in alpha):
             torch._foreach_add_(params, torch._foreach_mul(update, alpha))
 
-    def _loss(self, step_size: float, vars: Vars, closure, params: list[torch.Tensor],
+    def _loss(self, step_size: float, var: Var, closure, params: list[torch.Tensor],
               update: list[torch.Tensor], backward:bool=False) -> float:
 
         # if step_size is 0, we might already know the loss
-        if (vars.loss is not None) and (step_size == 0):
-            return tofloat(vars.loss)
+        if (var.loss is not None) and (step_size == 0):
+            return tofloat(var.loss)
 
         # check max iter
         if self._maxiter is not None and self._current_iter >= self._maxiter: raise MaxLineSearchItersReached
@@ -85,23 +85,23 @@ class LineSearch(Module, ABC):
             self._lowest_loss = tofloat(loss)
             self._best_step_size = step_size
 
-        # if evaluated loss at step size 0, set it to vars.loss
+        # if evaluated loss at step size 0, set it to var.loss
         if step_size == 0:
-            vars.loss = loss
-            if backward: vars.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+            var.loss = loss
+            if backward: var.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
 
         return tofloat(loss)
 
-    def _loss_derivative(self, step_size: float, vars: Vars, closure,
+    def _loss_derivative(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 (vars.grad is not None) and (step_size == 0):
-            loss = self._loss(step_size=step_size,vars=vars,closure=closure,params=params,update=update,backward=False)
-            derivative = - sum(t.sum() for t in torch._foreach_mul(vars.grad, update))
+        if (var.grad is not None) and (step_size == 0):
+            loss = self._loss(step_size=step_size,var=var,closure=closure,params=params,update=update,backward=False)
+            derivative = - sum(t.sum() for t in torch._foreach_mul(var.grad, update))
 
         else:
             # loss with a backward pass sets params.grad
-            loss = self._loss(step_size=step_size,vars=vars,closure=closure,params=params,update=update,backward=True)
+            loss = self._loss(step_size=step_size,var=var,closure=closure,params=params,update=update,backward=True)
 
             # directional derivative
             derivative = - sum(t.sum() for t in torch._foreach_mul([p.grad if p.grad is not None
@@ -109,60 +109,60 @@ class LineSearch(Module, ABC):
 
         return loss, tofloat(derivative)
 
-    def evaluate_step_size(self, step_size: float, vars: Vars, backward:bool=False):
-        closure = vars.closure
+    def evaluate_step_size(self, step_size: float, var: Var, backward:bool=False):
+        closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return self._loss(step_size=step_size, vars=vars, closure=closure, params=vars.params,update=vars.get_update(),backward=backward)
+        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, vars: Vars):
-        closure = vars.closure
+    def evaluate_step_size_loss_and_derivative(self, step_size: float, var: Var):
+        closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return self._loss_derivative(step_size=step_size, vars=vars, closure=closure, params=vars.params,update=vars.get_update())
+        return self._loss_derivative(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_update())
 
-    def make_objective(self, vars: Vars, backward:bool=False):
-        closure = vars.closure
+    def make_objective(self, var: Var, backward:bool=False):
+        closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return partial(self._loss, vars=vars, closure=closure, params=vars.params, update=vars.get_update(), backward=backward)
+        return partial(self._loss, var=var, closure=closure, params=var.params, update=var.get_update(), backward=backward)
 
-    def make_objective_with_derivative(self, vars: Vars):
-        closure = vars.closure
+    def make_objective_with_derivative(self, var: Var):
+        closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return partial(self._loss_derivative, vars=vars, closure=closure, params=vars.params, update=vars.get_update())
+        return partial(self._loss_derivative, var=var, closure=closure, params=var.params, update=var.get_update())
 
     @abstractmethod
-    def search(self, update: list[torch.Tensor], vars: Vars) -> float:
+    def search(self, update: list[torch.Tensor], var: Var) -> float:
         """Finds the step size to use"""
 
     @torch.no_grad
-    def step(self, vars: Vars) -> Vars:
+    def step(self, var: Var) -> Var:
         self._reset()
-        params = vars.params
-        update = vars.get_update()
+        params = var.params
+        update = var.get_update()
 
         try:
-            step_size = self.search(update=update, vars=vars)
+            step_size = self.search(update=update, var=var)
         except MaxLineSearchItersReached:
             step_size = self._best_step_size
 
         # set loss_approx
-        if vars.loss_approx is None: vars.loss_approx = self._lowest_loss
+        if var.loss_approx is None: var.loss_approx = self._lowest_loss
 
         # this is last module - set step size to found step_size times lr
-        if vars.is_last:
+        if var.is_last:
 
-            if vars.last_module_lrs is None:
+            if var.last_module_lrs is None:
                 self.set_step_size_(step_size, params=params, update=update)
 
             else:
-                self._set_per_parameter_step_size_([step_size*lr for lr in vars.last_module_lrs], params=params, update=update)
+                self._set_per_parameter_step_size_([step_size*lr for lr in var.last_module_lrs], params=params, update=update)
 
-            vars.stop = True; vars.skip_update = True
-            return vars
+            var.stop = True; var.skip_update = True
+            return var
 
         # revert parameters and multiply update by step size
         self.set_step_size_(0, params=params, update=update)
-        torch._foreach_mul_(vars.update, step_size)
-        return vars
+        torch._foreach_mul_(var.update, step_size)
+        return var
 
 
 class GridLineSearch(LineSearch):
@@ -172,10 +172,10 @@ class GridLineSearch(LineSearch):
         super().__init__(defaults)
 
     @torch.no_grad
-    def search(self, update, vars):
-        start,end,num=itemgetter('start','end','num')(self.settings[vars.params[0]])
+    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(), vars=vars, backward=False)
+            self.evaluate_step_size(lr.item(), var=var, backward=False)
 
         return self._best_step_size

torchzero/modules/line_search/scipy.py

@@ -7,6 +7,21 @@ from .line_search import LineSearch
 
 
 class ScipyMinimizeScalar(LineSearch):
+    """Line search via :code:`scipy.optimize.minimize_scalar` which implements brent, golden search and bounded brent methods.
+
+    Args:
+        method (str | None, optional): "brent", "golden" or "bounded". Defaults to None.
+        maxiter (int | None, optional): maximum number of function evaluations the line search is allowed to perform. Defaults to None.
+        bracket (Sequence | None, optional):
+            Either a triple (xa, xb, xc) satisfying xa < xb < xc and func(xb) < func(xa) and  func(xb) < func(xc), or a pair (xa, xb) to be used as initial points for a downhill bracket search. Defaults to None.
+        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.
+        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
+
+    """
     def __init__(
         self,
         method: str | None = None,
@@ -24,10 +39,10 @@ class ScipyMinimizeScalar(LineSearch):
 
 
     @torch.no_grad
-    def search(self, update, vars):
-        objective = self.make_objective(vars=vars)
+    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[vars.params[0]])
+            'method', 'bracket', 'bounds', 'tol', 'options', 'maxiter')(self.settings[var.params[0]])
 
         if maxiter is not None:
             options = dict(options) if isinstance(options, Mapping) else {}

torchzero/modules/line_search/strong_wolfe.py

@@ -183,6 +183,21 @@ def _notfinite(x):
     return not math.isfinite(x)
 
 class StrongWolfe(LineSearch):
+    """Cubic interpolation line search satisfying Strong Wolfe condition.
+
+    Args:
+        init (float, optional): Initial step size. Defaults to 1.0.
+        c1 (float, optional): Acceptance value for weak wolfe condition. Defaults to 1e-4.
+        c2 (float, optional): Acceptance value for strong wolfe condition (set to 0.1 for conjugate gradient). Defaults to 0.9.
+        maxiter (int, optional): Maximum number of line search iterations. Defaults to 25.
+        maxzoom (int, optional): Maximum number of zoom iterations. Defaults to 10.
+        expand (float, optional): Expansion factor (multipler to step size when weak condition not satisfied). Defaults to 2.0.
+        adaptive (bool, optional):
+            when enabled, if line search failed, initial step size is reduced.
+            Otherwise it is reset to initial value. Defaults to True.
+        plus_minus (bool, optional):
+            If enabled and the direction is not descent direction, performs line search in opposite direction. Defaults to False.
+    """
     def __init__(
         self,
         init: float = 1.0,
@@ -193,23 +208,22 @@ class StrongWolfe(LineSearch):
         # a_max: float = 1e10,
         expand: float = 2.0,
         adaptive = True,
-        fallback = False,
         plus_minus = False,
     ):
         defaults=dict(init=init,c1=c1,c2=c2,maxiter=maxiter,maxzoom=maxzoom,
-                      expand=expand, adaptive=adaptive, fallback=fallback, plus_minus=plus_minus)
+                      expand=expand, adaptive=adaptive, plus_minus=plus_minus)
         super().__init__(defaults=defaults)
 
         self.global_state['initial_scale'] = 1.0
         self.global_state['beta_scale'] = 1.0
 
     @torch.no_grad
-    def search(self, update, vars):
-        objective = self.make_objective_with_derivative(vars=vars)
+    def search(self, update, var):
+        objective = self.make_objective_with_derivative(var=var)
 
-        init, c1, c2, maxiter, maxzoom, expand, adaptive, fallback, plus_minus = itemgetter(
+        init, c1, c2, maxiter, maxzoom, expand, adaptive, plus_minus = itemgetter(
             'init', 'c1', 'c2', 'maxiter', 'maxzoom',
-            'expand', 'adaptive', 'fallback', 'plus_minus')(self.settings[vars.params[0]])
+            'expand', 'adaptive', 'plus_minus')(self.settings[var.params[0]])
 
         f_0, g_0 = objective(0)
 
@@ -232,29 +246,4 @@ class StrongWolfe(LineSearch):
 
         # fallback to backtracking on fail
         if adaptive: self.global_state['initial_scale'] *= 0.5
-        if not fallback: return 0
-
-        objective = self.make_objective(vars=vars)
-
-        # # directional derivative
-        g_0 = -sum(t.sum() for t in torch._foreach_mul(vars.get_grad(), vars.get_update()))
-
-        step_size = backtracking_line_search(
-            objective,
-            g_0,
-            init=init * self.global_state["initial_scale"],
-            beta=0.5 * self.global_state["beta_scale"],
-            c=c1,
-            maxiter=maxiter * 2,
-            a_min=None,
-            try_negative=plus_minus,
-        )
-
-        # found an alpha that reduces loss
-        if step_size is not None:
-            self.global_state['beta_scale'] = min(1.0, self.global_state.get('beta_scale', 1) * math.sqrt(1.5))
-            return step_size
-
-        # on fail reduce beta scale value
-        self.global_state['beta_scale'] /= 1.5
-        return 0
+        return 0

torchzero/modules/line_search/trust_region.py

@@ -6,31 +6,43 @@ from .line_search import LineSearch
 
 
 class TrustRegion(LineSearch):
-    """Basic first order trust region, re-evaluates closure with updated parameters and scales step size based on function value change"""
+    """Basic first order trust region method. Re-evaluates the function after stepping, if value decreased sufficiently,
+    step size is increased. If value increased, step size is decreased. This is prone to collapsing.
+
+    Args:
+        nplus (float, optional): multiplier to step size on successful steps. Defaults to 1.5.
+        nminus (float, optional): multiplier to step size on unsuccessful steps. Defaults to 0.75.
+        c (float, optional): descent condition. Defaults to 1e-4.
+        init (float, optional): initial step size. Defaults to 1.
+        backtrack (bool, optional): whether to undo the step if value increased. Defaults to True.
+        adaptive (bool, optional):
+            If enabled, when multiple consecutive steps have been successful or unsuccessful,
+            the corresponding multipliers are increased, otherwise they are reset. Defaults to True.
+    """
     def __init__(self, nplus: float=1.5, nminus: float=0.75, c: float=1e-4, init: float = 1, backtrack: bool = True, adaptive: bool = True):
         defaults = dict(nplus=nplus, nminus=nminus, c=c, init=init, backtrack=backtrack, adaptive=adaptive)
         super().__init__(defaults)
 
     @torch.no_grad
-    def search(self, update, vars):
+    def search(self, update, var):
 
-        nplus, nminus, c, init, backtrack, adaptive = itemgetter('nplus','nminus','c','init','backtrack', 'adaptive')(self.settings[vars.params[0]])
+        nplus, nminus, c, init, backtrack, adaptive = itemgetter('nplus','nminus','c','init','backtrack', 'adaptive')(self.settings[var.params[0]])
         step_size = self.global_state.setdefault('step_size', init)
         previous_success = self.global_state.setdefault('previous_success', False)
         nplus_mul =  self.global_state.setdefault('nplus_mul', 1)
         nminus_mul = self.global_state.setdefault('nminus_mul', 1)
 
 
-        f_0 = self.evaluate_step_size(0, vars, backward=False)
+        f_0 = self.evaluate_step_size(0, var, backward=False)
 
         # directional derivative (0 if c = 0 because it is not needed)
         if c == 0: d = 0
-        else: d = -sum(t.sum() for t in torch._foreach_mul(vars.get_grad(), update))
+        else: d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), update))
 
         # test step size
         sufficient_f = f_0 + c * step_size * min(d, 0) # pyright:ignore[reportArgumentType]
 
-        f_1 = self.evaluate_step_size(step_size, vars, backward=False)
+        f_1 = self.evaluate_step_size(step_size, var, backward=False)
 
         proposed = step_size
 

torchzero/modules/lr/__init__.py

@@ -1,2 +1,2 @@
 from .lr import LR, StepSize, Warmup
-from .step_size import PolyakStepSize, RandomStepSize
+from .adaptive import PolyakStepSize, RandomStepSize

torchzero/modules/lr/{step_size.py → adaptive.py}

@@ -1,18 +1,20 @@
+"""Various step size strategies"""
 import random
 from typing import Any
-
+from operator import itemgetter
 import torch
 
 from ...core import Transform
-from ...utils import TensorList, NumberList
+from ...utils import TensorList, NumberList, unpack_dicts
 
 
 class PolyakStepSize(Transform):
-    """Polyak step-size.
+    """Polyak's step-size method.
 
     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.
+        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.
@@ -28,29 +30,24 @@ class PolyakStepSize(Transform):
         super().__init__(defaults, uses_grad=use_grad)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        loss = vars.get_loss(False)
+    def apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         tensors = TensorList(tensors)
         grads = TensorList(grads)
-        alpha = self.get_settings('alpha', params=params, cls=NumberList)
-        settings = self.settings[params[0]]
-        parameterwise = settings['parameterwise']
-        use_grad = settings['use_grad']
-        max = settings['max']
-        min_obj_value = settings['min_obj_value']
+        alpha = NumberList(s['alpha'] for s in settings)
+
+        parameterwise, use_grad, max, min_obj_value = itemgetter('parameterwise', 'use_grad', 'max', 'min_obj_value')(settings[0])
+
+        if use_grad: denom = tensors.dot(grads)
+        else: denom = tensors.dot(tensors)
 
         if parameterwise:
-            if use_grad: denom = (tensors * grads).sum()
-            else: denom = tensors.pow(2).sum()
             polyak_step_size: TensorList | Any = (loss - min_obj_value) / denom.where(denom!=0, 1)
             polyak_step_size = polyak_step_size.where(denom != 0, 0)
             if max is not None: polyak_step_size = polyak_step_size.clamp_max(max)
 
         else:
-            if use_grad: denom = tensors.dot(grads)
-            else: denom = tensors.dot(tensors)
-            if denom == 0: polyak_step_size = 0 # we converged
+            if denom.abs() <= torch.finfo(denom.dtype).eps: polyak_step_size = 0 # converged
             else: polyak_step_size = (loss - min_obj_value) / denom
 
             if max is not None:
@@ -60,9 +57,8 @@ class PolyakStepSize(Transform):
         return tensors
 
 
-
 class RandomStepSize(Transform):
-    """Uses random global step size from `low` to `high`.
+    """Uses random global or layer-wise step size from `low` to `high`.
 
     Args:
         low (float, optional): minimum learning rate. Defaults to 0.
@@ -76,21 +72,21 @@ class RandomStepSize(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        settings = self.settings[params[0]]
-        parameterwise = settings['parameterwise']
+    def apply(self, tensors, params, grads, loss, states, settings):
+        s = settings[0]
+        parameterwise = s['parameterwise']
 
-        seed = settings['seed']
+        seed = s['seed']
         if 'generator' not in self.global_state:
             self.global_state['generator'] = random.Random(seed)
         generator: random.Random = self.global_state['generator']
 
         if parameterwise:
-            low, high = self.get_settings('low', 'high', params=params)
+            low, high = unpack_dicts(settings, 'low', 'high')
             lr = [generator.uniform(l, h) for l, h in zip(low, high)]
         else:
-            low = settings['low']
-            high = settings['high']
+            low = s['low']
+            high = s['high']
             lr = generator.uniform(low, high)
 
         torch._foreach_mul_(tensors, lr)

torchzero/modules/lr/lr.py

@@ -1,8 +1,8 @@
+"""Learning rate"""
 import torch
 
 from ...core import Transform
-from ...utils import NumberList, TensorList, generic_eq
-
+from ...utils import NumberList, TensorList, generic_eq, unpack_dicts
 
 def lazy_lr(tensors: TensorList, lr: float | list, inplace:bool):
     """multiplies by lr if lr is not 1"""
@@ -11,48 +11,52 @@ def lazy_lr(tensors: TensorList, lr: float | list, inplace:bool):
     return tensors * lr
 
 class LR(Transform):
+    """Learning rate. Adding this module also adds support for LR schedulers."""
     def __init__(self, lr: float):
         defaults=dict(lr=lr)
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        return lazy_lr(TensorList(tensors), lr=self.get_settings('lr', params=params), inplace=True)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        return lazy_lr(TensorList(tensors), lr=[s['lr'] for s in settings], inplace=True)
 
 class StepSize(Transform):
-    """this is exactly the same as LR, except the `lr` parameter can be renamed to any other name"""
+    """this is exactly the same as LR, except the `lr` parameter can be renamed to any other name to avoid clashes"""
     def __init__(self, step_size: float, key = 'step_size'):
         defaults={"key": key, key: step_size}
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        lrs = []
-        for p in params:
-            settings = self.settings[p]
-            lrs.append(settings[settings['key']])
-        return lazy_lr(TensorList(tensors), lr=lrs, inplace=True)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        return lazy_lr(TensorList(tensors), lr=[s[s['key']] for s in settings], inplace=True)
 
 
-def warmup(step: int, start_lr: float | NumberList, end_lr: float | NumberList, steps: float):
+def _warmup_lr(step: int, start_lr: float | NumberList, end_lr: float | NumberList, steps: float):
     """returns warm up lr scalar"""
     if step > steps: return end_lr
     return start_lr + (end_lr - start_lr) * (step / steps)
 
 class Warmup(Transform):
+    """Learning rate warmup, linearly increases learning rate multiplier from :code:`start_lr` to :code:`end_lr` over :code:`steps` steps.
+
+    Args:
+        start_lr (_type_, optional): initial learning rate multiplier on first step. Defaults to 1e-5.
+        end_lr (float, optional): learning rate multiplier at the end and after warmup. Defaults to 1.
+        steps (int, optional): number of steps to perform warmup for. Defaults to 100.
+    """
     def __init__(self, start_lr = 1e-5, end_lr:float = 1, steps = 100):
         defaults = dict(start_lr=start_lr,end_lr=end_lr, steps=steps)
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        start_lr, end_lr = self.get_settings('start_lr', 'end_lr', params=params, cls = NumberList)
-        num_steps = self.settings[params[0]]['steps']
+    def apply(self, tensors, params, grads, loss, states, settings):
+        start_lr, end_lr = unpack_dicts(settings, 'start_lr', 'end_lr', cls = NumberList)
+        num_steps = settings[0]['steps']
         step = self.global_state.get('step', 0)
 
         target = lazy_lr(
             TensorList(tensors),
-            lr=warmup(step=step, start_lr=start_lr, end_lr=end_lr, steps=num_steps),
+            lr=_warmup_lr(step=step, start_lr=start_lr, end_lr=end_lr, steps=num_steps),
             inplace=True
         )
         self.global_state['step'] = step + 1

torchzero/modules/momentum/averaging.py

@@ -1,3 +1,4 @@
+"""Modules that perform averaging over a history of past updates."""
 from collections import deque
 from collections.abc import Sequence
 from typing import Any, Literal, cast
@@ -9,14 +10,19 @@ from ...utils import tolist
 
 
 class Averaging(TensorwiseTransform):
+    """Average of past :code:`history_size` updates.
+
+    Args:
+        history_size (int): Number of past updates to average
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, history_size: int, target: Target = 'update'):
         defaults = dict(history_size=history_size)
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensor, param, grad, vars):
-        history_size = self.settings[param]['history_size']
-        state = self.state[param]
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+        history_size = settings['history_size']
         if 'history' not in state:
             state['history'] = deque(maxlen=history_size)
             state['average'] = torch.zeros_like(tensor)
@@ -29,15 +35,19 @@ class Averaging(TensorwiseTransform):
         return average / len(history)
 
 class WeightedAveraging(TensorwiseTransform):
-    """weights are oldest to newest"""
+    """Weighted average of past :code:`len(weights)` updates.
+
+    Args:
+        weights (Sequence[float]): a sequence of weights from oldest to newest.
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, weights: Sequence[float] | torch.Tensor | Any, target: Target = 'update'):
         defaults = dict(weights = tolist(weights))
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensor, param, grad, vars):
-        weights = self.settings[param]['weights']
-        state = self.state[param]
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+        weights = settings['weights']
 
         if 'history' not in state:
             state['history'] = deque(maxlen=len(weights))
@@ -59,14 +69,19 @@ class WeightedAveraging(TensorwiseTransform):
 
 
 class MedianAveraging(TensorwiseTransform):
+    """Median of past :code:`history_size` updates.
+
+    Args:
+        history_size (int): Number of past updates to average
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, history_size: int, target: Target = 'update'):
         defaults = dict(history_size = history_size)
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensor, param, grad, vars):
-        history_size = self.settings[param]['history_size']
-        state = self.state[param]
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+        history_size = settings['history_size']
 
         if 'history' not in state:
             state['history'] = deque(maxlen=history_size)