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

torchzero/modules/momentum/cautious.py

@@ -1,3 +1,4 @@
+"""Cautioning related modules"""
 from collections import deque
 from operator import itemgetter
 from typing import Literal
@@ -5,7 +6,7 @@ from typing import Literal
 import torch
 
 from ...core import Target, Transform, Module, Chainable
-from ...utils import NumberList, TensorList
+from ...utils import NumberList, TensorList, unpack_dicts
 
 
 def cautious_(
@@ -64,27 +65,33 @@ class Cautious(Transform):
         normalize=False,
         eps=1e-6,
         mode: Literal["zero", "grad", "backtrack"] = "zero",
-        target: Target = "update",
     ):
         defaults = dict(normalize=normalize, eps=eps, mode=mode)
-        super().__init__(defaults, uses_grad=True, target=target)
+        super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
-        mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(self.settings[params[0]])
+        mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(settings[0])
         return cautious_(TensorList(tensors), TensorList(grads), normalize=normalize, eps=eps, mode=mode)
 
 class UpdateGradientSignConsistency(Transform):
-    """1 where signs match 0 otherwise"""
-    def __init__(self, normalize = False, eps=1e-6, target: Target = 'update'):
+    """Compares update and gradient signs. Output will have 1s where signs match, and 0s where they don't.
+
+    Args:
+        normalize (bool, optional):
+            renormalize update after masking. Defaults to False.
+        eps (float, optional): epsilon for normalization. Defaults to 1e-6.
+    """
+    def __init__(self, normalize = False, eps=1e-6):
+
         defaults = dict(normalize=normalize, eps=eps)
-        super().__init__(defaults, uses_grad=True, target=target)
+        super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
-        normalize, eps = itemgetter('normalize', 'eps')(self.settings[params[0]])
+        normalize, eps = itemgetter('normalize', 'eps')(settings[0])
 
         mask = (TensorList(tensors).mul_(grads)).gt_(0)
         if normalize: mask = mask / mask.global_mean().clip(min = eps) # pyright: ignore[reportOperatorIssue]
@@ -92,6 +99,23 @@ class UpdateGradientSignConsistency(Transform):
         return mask
 
 class IntermoduleCautious(Module):
+    """Negaties update on :code:`main` module where it's sign doesn't match with output of :code:`compare` module.
+
+    Args:
+        main (Chainable): main module or sequence of modules whose update will be cautioned.
+        compare (Chainable): modules or sequence of modules to compare the sign to.
+        normalize (bool, optional):
+            renormalize update after masking. Defaults to False.
+        eps (float, optional): epsilon for normalization. Defaults to 1e-6.
+        mode (str, optional):
+            what to do with updates with inconsistent signs.
+
+            "zero" - set them to zero (as in paper)
+
+            "grad" - set them to the gradient
+
+            "backtrack" - negate them (same as using update magnitude and gradient sign)
+    """
     def __init__(
         self,
         main: Chainable,
@@ -100,6 +124,7 @@ class IntermoduleCautious(Module):
         eps=1e-6,
         mode: Literal["zero", "grad", "backtrack"] = "zero",
     ):
+
         defaults = dict(normalize=normalize, eps=eps, mode=mode)
         super().__init__(defaults)
 
@@ -107,40 +132,45 @@ class IntermoduleCautious(Module):
         self.set_child('compare', compare)
 
     @torch.no_grad
-    def step(self, vars):
+    def step(self, var):
         main = self.children['main']
         compare = self.children['compare']
 
-        main_vars = main.step(vars.clone(clone_update=True))
-        vars.update_attrs_from_clone_(main_vars)
+        main_var = main.step(var.clone(clone_update=True))
+        var.update_attrs_from_clone_(main_var)
 
-        compare_vars = compare.step(vars.clone(clone_update=True))
-        vars.update_attrs_from_clone_(compare_vars)
+        compare_var = compare.step(var.clone(clone_update=True))
+        var.update_attrs_from_clone_(compare_var)
 
-        mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(self.settings[vars.params[0]])
-        vars.update = cautious_(
-            TensorList(main_vars.get_update()),
-            TensorList(compare_vars.get_update()),
+        mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(self.settings[var.params[0]])
+        var.update = cautious_(
+            TensorList(main_var.get_update()),
+            TensorList(compare_var.get_update()),
             normalize=normalize,
             mode=mode,
             eps=eps,
         )
 
-        return vars
+        return var
 
 class ScaleByGradCosineSimilarity(Transform):
+    """Multiplies the update by cosine similarity with gradient.
+    If cosine similarity is negative, naturally the update will be negated as well.
+
+    Args:
+        eps (float, optional): epsilon for division. Defaults to 1e-6.
+    """
     def __init__(
         self,
-        eps=1e-6,
-        target: Target = "update",
+        eps: float = 1e-6,
     ):
         defaults = dict(eps=eps)
-        super().__init__(defaults, uses_grad=True, target=target)
+        super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
-        eps = self.settings[params[0]]['eps']
+        eps = settings[0]['eps']
         tensors = TensorList(tensors)
         grads = TensorList(grads)
         cos_sim = (tensors.dot(grads)) / (tensors.global_vector_norm() * grads.global_vector_norm()).clip(min=eps)
@@ -148,6 +178,14 @@ class ScaleByGradCosineSimilarity(Transform):
         return tensors.mul_(cos_sim)
 
 class ScaleModulesByCosineSimilarity(Module):
+    """Scales the output of :code:`main` module by it's cosine similarity to the output
+    of :code:`compare` module.
+
+    Args:
+        main (Chainable): main module or sequence of modules whose update will be scaled.
+        compare (Chainable): module or sequence of modules to compare to
+        eps (float, optional): epsilon for division. Defaults to 1e-6.
+    """
     def __init__(
         self,
         main: Chainable,
@@ -161,21 +199,21 @@ class ScaleModulesByCosineSimilarity(Module):
         self.set_child('compare', compare)
 
     @torch.no_grad
-    def step(self, vars):
+    def step(self, var):
         main = self.children['main']
         compare = self.children['compare']
 
-        main_vars = main.step(vars.clone(clone_update=True))
-        vars.update_attrs_from_clone_(main_vars)
+        main_var = main.step(var.clone(clone_update=True))
+        var.update_attrs_from_clone_(main_var)
 
-        compare_vars = compare.step(vars.clone(clone_update=True))
-        vars.update_attrs_from_clone_(compare_vars)
+        compare_var = compare.step(var.clone(clone_update=True))
+        var.update_attrs_from_clone_(compare_var)
 
-        m = TensorList(main_vars.get_update())
-        c = TensorList(compare_vars.get_update())
-        eps = self.settings[vars.params[0]]['eps']
+        m = TensorList(main_var.get_update())
+        c = TensorList(compare_var.get_update())
+        eps = self.settings[var.params[0]]['eps']
 
         cos_sim = (m.dot(c)) / (m.global_vector_norm() * c.global_vector_norm()).clip(min=eps)
 
-        vars.update = m.mul_(cos_sim)
-        return vars
+        var.update = m.mul_(cos_sim)
+        return var

torchzero/modules/momentum/ema.py

@@ -5,18 +5,19 @@ from typing import Literal
 import torch
 
 from ...core import Target, Transform
-from ...utils import TensorList, NumberList
+from ...utils import TensorList, NumberList, unpack_dicts, unpack_states
 from ..functional import debias, ema_, ema_sq_, sqrt_ema_sq_, centered_ema_sq_, sqrt_centered_ema_sq_, debias_second_momentum
 
 
 class EMA(Transform):
-    """Maintains EMA of update.
+    """Maintains an exponential moving average of update.
 
     Args:
         momentum (float, optional): momentum (beta). Defaults to 0.9.
         dampening (float, optional): momentum dampening. Defaults to 0.
         debiased (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
         lerp (bool, optional): whether to use linear interpolation. Defaults to True.
+        ema_init (str, optional): initial values for the EMA, "zeros" or "update".
         target (Target, optional): target to apply EMA to. Defaults to 'update'.
     """
     def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=True, ema_init: Literal['zeros', 'update'] = 'zeros', target: Target = 'update'):
@@ -24,13 +25,14 @@ class EMA(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        debiased, lerp, ema_init = itemgetter('debiased','lerp','ema_init')(self.settings[params[0]])
+        debiased, lerp, ema_init = itemgetter('debiased','lerp','ema_init')(settings[0])
 
-        exp_avg = self.get_state('exp_avg', params=params, init=torch.zeros_like if ema_init=='zeros' else tensors, cls=TensorList)
-        momentum, dampening = self.get_settings('momentum','dampening', params=params, cls=NumberList)
+        exp_avg = unpack_states(states, tensors, 'exp_avg',
+                                init=torch.zeros_like if ema_init=='zeros' else tensors, cls=TensorList)
+        momentum, dampening = unpack_dicts(settings, 'momentum','dampening', cls=NumberList)
 
         exp_avg = ema_(TensorList(tensors), exp_avg_=exp_avg,beta=momentum,dampening=dampening,lerp=lerp)
 
@@ -39,44 +41,58 @@ class EMA(Transform):
 
 
 class EMASquared(Transform):
+    """Maintains an exponential moving average of squared updates.
+
+    Args:
+        beta (float, optional): momentum value. Defaults to 0.999.
+        amsgrad (bool, optional): whether to maintain maximum of the exponential moving average. Defaults to False.
+        pow (float, optional): power, absolute value is always used. Defaults to 2.
+    """
     EMA_SQ_FN: staticmethod = staticmethod(ema_sq_)
 
-    def __init__(self, beta:float=0.999, amsgrad=False, pow:float=2, target: Target = 'update'):
+    def __init__(self, beta:float=0.999, amsgrad=False, pow:float=2):
         defaults = dict(beta=beta,pow=pow,amsgrad=amsgrad)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         amsgrad, pow = itemgetter('amsgrad', 'pow')(self.settings[params[0]])
-        beta = self.get_settings('beta', params=params, cls=NumberList)
+        beta = NumberList(s['beta'] for s in settings)
 
         if amsgrad:
-            exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+            exp_avg_sq, max_exp_avg_sq = unpack_states(states, tensors, 'exp_avg_sq', 'max_exp_avg_sq', cls=TensorList)
         else:
-            exp_avg_sq = self.get_state('exp_avg_sq', params=params, cls=TensorList)
+            exp_avg_sq = unpack_states(states, tensors, 'exp_avg_sq', cls=TensorList)
             max_exp_avg_sq = None
 
         return self.EMA_SQ_FN(TensorList(tensors), exp_avg_sq_=exp_avg_sq, beta=beta, max_exp_avg_sq_=max_exp_avg_sq, pow=pow).clone()
 
 class SqrtEMASquared(Transform):
-    SQRT_EMA_SQ_FN: staticmethod = staticmethod(sqrt_ema_sq_)
+    """Maintains an exponential moving average of squared updates, outputs optionally debiased square root.
 
-    def __init__(self, beta:float=0.999, amsgrad=False, debiased: bool = False, pow:float=2, target: Target = 'update',):
+    Args:
+        beta (float, optional): momentum value. Defaults to 0.999.
+        amsgrad (bool, optional): whether to maintain maximum of the exponential moving average. Defaults to False.
+        debiased (bool, optional): whether to multiply the output by a debiasing term from the Adam method. Defaults to False.
+        pow (float, optional): power, absolute value is always used. Defaults to 2.
+    """
+    SQRT_EMA_SQ_FN: staticmethod = staticmethod(sqrt_ema_sq_)
+    def __init__(self, beta:float=0.999, amsgrad=False, debiased: bool = False, pow:float=2,):
         defaults = dict(beta=beta,pow=pow,amsgrad=amsgrad,debiased=debiased)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, uses_grad=False)
 
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        amsgrad, pow, debiased = itemgetter('amsgrad', 'pow', 'debiased')(self.settings[params[0]])
-        beta = self.get_settings('beta', params=params, cls=NumberList)
+        amsgrad, pow, debiased = itemgetter('amsgrad', 'pow', 'debiased')(settings[0])
+        beta = NumberList(s['beta'] for s in settings)
 
         if amsgrad:
-            exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+            exp_avg_sq, max_exp_avg_sq = unpack_states(states, tensors, 'exp_avg_sq', 'max_exp_avg_sq', cls=TensorList)
         else:
-            exp_avg_sq = self.get_state('exp_avg_sq', params=params, cls=TensorList)
+            exp_avg_sq = unpack_states(states, tensors, 'exp_avg_sq', cls=TensorList)
             max_exp_avg_sq = None
 
         return self.SQRT_EMA_SQ_FN(
@@ -91,47 +107,73 @@ class SqrtEMASquared(Transform):
 
 
 class Debias(Transform):
+    """Multiplies the update by an Adam debiasing term based first and/or second momentum.
+
+    Args:
+        beta1 (float | None, optional):
+            first momentum, should be the same as first momentum used in modules before. Defaults to None.
+        beta2 (float | None, optional):
+            second (squared) momentum, should be the same as second momentum used in modules before. Defaults to None.
+        alpha (float, optional): learning rate. Defaults to 1.
+        pow (float, optional): power, assumes absolute value is used. Defaults to 2.
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, beta1: float | None = None, beta2: float | None = None, alpha: float = 1, pow:float=2, target: Target = 'update',):
         defaults = dict(beta1=beta1, beta2=beta2, alpha=alpha, pow=pow)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        settings = self.settings[params[0]]
-        pow = settings['pow']
-        alpha, beta1, beta2 = self.get_settings('alpha', 'beta1', 'beta2', params=params, cls=NumberList)
+        pow = settings[0]['pow']
+        alpha, beta1, beta2 = unpack_dicts(settings, 'alpha', 'beta1', 'beta2', cls=NumberList)
 
         return debias(TensorList(tensors), step=step, beta1=beta1, beta2=beta2, alpha=alpha, pow=pow, inplace=True)
 
 class Debias2(Transform):
+    """Multiplies the update by an Adam debiasing term based on the second momentum.
+
+    Args:
+        beta (float | None, optional):
+            second (squared) momentum, should be the same as second momentum used in modules before. Defaults to None.
+        pow (float, optional): power, assumes absolute value is used. Defaults to 2.
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, beta: float = 0.999, pow: float = 2, target: Target = 'update',):
         defaults = dict(beta=beta, pow=pow)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        pow = self.settings[params[0]]['pow']
-        beta = self.get_settings('beta', params=params, cls=NumberList)
+        pow = settings[0]['pow']
+        beta = NumberList(s['beta'] for s in settings)
         return debias_second_momentum(TensorList(tensors), step=step, beta=beta, pow=pow, inplace=True)
 
 class CenteredEMASquared(Transform):
-    def __init__(self, beta: float = 0.99, amsgrad=False, pow:float=2, target: Target = 'update'):
+    """Maintains a centered exponential moving average of squared updates. This also maintains an additional
+    exponential moving average of un-squared updates, square of which is subtracted from the EMA.
+
+    Args:
+        beta (float, optional): momentum value. Defaults to 0.999.
+        amsgrad (bool, optional): whether to maintain maximum of the exponential moving average. Defaults to False.
+        pow (float, optional): power, absolute value is always used. Defaults to 2.
+    """
+    def __init__(self, beta: float = 0.99, amsgrad=False, pow:float=2):
         defaults = dict(beta=beta, amsgrad=amsgrad, pow=pow)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        amsgrad, pow = itemgetter('amsgrad', 'pow')(self.settings[params[0]])
-        beta = self.get_settings('beta', params=params, cls=NumberList)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        amsgrad, pow = itemgetter('amsgrad', 'pow')(settings[0])
+        beta = NumberList(s['beta'] for s in settings)
 
         if amsgrad:
-            exp_avg, exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+            exp_avg, exp_avg_sq, max_exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', cls=TensorList)
         else:
-            exp_avg, exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', params=params, cls=TensorList)
+            exp_avg, exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', cls=TensorList)
             max_exp_avg_sq = None
 
         return centered_ema_sq_(
@@ -144,21 +186,30 @@ class CenteredEMASquared(Transform):
         ).clone()
 
 class CenteredSqrtEMASquared(Transform):
-    def __init__(self, beta: float = 0.99, amsgrad=False, debiased: bool = False, pow:float=2, target: Target = 'update'):
+    """Maintains a centered exponential moving average of squared updates, outputs optionally debiased square root.
+    This also maintains an additional exponential moving average of un-squared updates, square of which is subtracted from the EMA.
+
+    Args:
+        beta (float, optional): momentum value. Defaults to 0.999.
+        amsgrad (bool, optional): whether to maintain maximum of the exponential moving average. Defaults to False.
+        debiased (bool, optional): whether to multiply the output by a debiasing term from the Adam method. Defaults to False.
+        pow (float, optional): power, absolute value is always used. Defaults to 2.
+    """
+    def __init__(self, beta: float = 0.99, amsgrad=False, debiased: bool = False, pow:float=2):
         defaults = dict(beta=beta, amsgrad=amsgrad, debiased=debiased, pow=pow)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        amsgrad, pow, debiased = itemgetter('amsgrad', 'pow', 'debiased')(self.settings[params[0]])
-        beta = self.get_settings('beta', params=params, cls=NumberList)
+        amsgrad, pow, debiased = itemgetter('amsgrad', 'pow', 'debiased')(settings[0])
+        beta = NumberList(s['beta'] for s in settings)
 
         if amsgrad:
-            exp_avg, exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+            exp_avg, exp_avg_sq, max_exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', cls=TensorList)
         else:
-            exp_avg, exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', params=params, cls=TensorList)
+            exp_avg, exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', cls=TensorList)
             max_exp_avg_sq = None
 
         return sqrt_centered_ema_sq_(

torchzero/modules/momentum/experimental.py

@@ -6,7 +6,7 @@ from typing import Literal
 import torch
 
 from ...core import Target, Transform
-from ...utils import NumberList, TensorList
+from ...utils import NumberList, TensorList, unpack_states, unpack_dicts
 from ..functional import ema_, ema_sq_, sqrt_ema_sq_
 from .ema import EMASquared, SqrtEMASquared
 from .momentum import nag_
@@ -43,22 +43,22 @@ def precentered_ema_sq_(
     return exp_avg_sq_
 
 class PrecenteredEMASquared(Transform):
+    """Maintains un-squared EMA, the updates are centered by it before being fed into squared EMA."""
     def __init__(self, beta1:float=0.99, beta2=0.99, min_step: int = 2, amsgrad=False, pow:float=2, target: Target = 'update'):
         defaults = dict(beta1=beta1,beta2=beta2,pow=pow,amsgrad=amsgrad, min_step=min_step)
         super().__init__(defaults, uses_grad=False, target=target)
-        self.current_step = 0
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        self.current_step += 1
+    def apply(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        beta1, beta2 = self.get_settings('beta1','beta2', params=params, cls=NumberList)
-        amsgrad, pow, min_step = itemgetter('amsgrad', 'pow', 'min_step')(self.settings[params[0]])
+        beta1, beta2 = unpack_dicts(settings, 'beta1','beta2', cls=NumberList)
+        amsgrad, pow, min_step = itemgetter('amsgrad', 'pow', 'min_step')(settings[0])
 
         if amsgrad:
-            exp_avg, exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', params=params, cls=TensorList)
+            exp_avg, exp_avg_sq, max_exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', cls=TensorList)
         else:
-            exp_avg, exp_avg_sq = self.get_state('exp_avg', 'exp_avg_sq', params=params, cls=TensorList)
+            exp_avg, exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', cls=TensorList)
             max_exp_avg_sq = None
 
         return precentered_ema_sq_(
@@ -67,7 +67,7 @@ class PrecenteredEMASquared(Transform):
             exp_avg_sq_=exp_avg_sq,
             beta1=beta1,
             beta2=beta2,
-            step = self.current_step,
+            step = step,
             min_step=min_step,
             pow=pow,
             max_exp_avg_sq_=max_exp_avg_sq,
@@ -119,9 +119,11 @@ def sqrt_nag_ema_sq_(
                         pow=pow,debiased=debiased,step=step,ema_sq_fn=partial(nag_ema_sq_,lerp=lerp))
 
 class NesterovEMASquared(EMASquared):
+    """squared momentum with nesterov momentum rule"""
     EMA_SQ_FN = staticmethod(nag_ema_sq_)
 
 class SqrtNesterovEMASquared(SqrtEMASquared):
+    """square root of squared momentum with nesterov momentum rule"""
     SQRT_EMA_SQ_FN = staticmethod(sqrt_nag_ema_sq_)
 
 
@@ -141,14 +143,20 @@ def coordinate_momentum_(
 
 
 class CoordinateMomentum(Transform):
+    """Maintains a momentum buffer, on each step each value in the buffer has :code:`p` chance to be updated with the new value.
+
+    Args:
+        p (float, optional): _description_. Defaults to 0.1.
+        target (Target, optional): _description_. Defaults to 'update'.
+    """
     def __init__(self, p: float = 0.1, target: Target = 'update'):
         defaults = dict(p=p)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        p = self.get_settings('p', params=params, cls=NumberList)
-        velocity = self.get_state('velocity', params=params, cls=TensorList)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        p = NumberList(s['p'] for s in settings)
+        velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
         return coordinate_momentum_(TensorList(tensors), velocity_=velocity, p=p).clone()
 
 
@@ -180,7 +188,7 @@ class CoordinateMomentum(Transform):
 #         super().__init__(defaults, uses_grad=False)
 
 #     @torch.no_grad
-#     def transform(self, tensors, params, grads, vars):
+#     def apply(self, tensors, params, grads, loss, states, settings):
 #         momentum,dampening = self.get_settings('momentum','dampening', params=params, cls=NumberList)
 #         abs,lerp,normalize_velocity = self.first_setting('abs','lerp','normalize_velocity', params=params)
 #         velocity = self.get_state('velocity', params=params, cls=TensorList)