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

torchzero/core/transform.py

@@ -1,18 +1,18 @@
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Sequence
-from typing import Any, Literal
+from collections.abc import Iterable, Sequence, Mapping
+from typing import Any, Literal, final
 
 import torch
 
-from ..utils import set_storage_
-from .module import Module, Vars, Chain, Chainable
+from ..utils import set_storage_, TensorList, vec_to_tensors
+from .module import Module, Var, Chain, Chainable
 
 Target = Literal['grad', 'update', 'closure', 'params_direct', 'params_difference', 'update_difference']
 
 class Transform(Module, ABC):
-    """Base class for a transform.
+    """Base class for a transform. This is an abstract class, to use it, subclass it and override `update` and `apply` methods.
 
-    This is an abstract class, to use it, subclass it and override `transform`.
+    A transform is a module that can also be applied manually to an arbitrary sequence of tensors.
 
     Args:
         defaults (dict[str,Any] | None): dict with default values.
@@ -20,62 +20,180 @@ class Transform(Module, ABC):
             Set this to True if `transform` method uses the `grad` argument. This will ensure
             `grad` is always computed and can't be None. Otherwise set to False.
         target (Target, optional):
-            what to set on vars. Defaults to 'update'.
+            what to set on var. Defaults to 'update'.
     """
-    def __init__(self, defaults: dict[str,Any] | None, uses_grad: bool, target: Target = 'update'):
+    def __init__(
+        self,
+        defaults: dict[str,Any] | None,
+        uses_grad: bool,
+        concat_params: bool = False,
+        update_freq: int = 1,
+        scale_first: bool = False,
+        inner: Chainable | None = None,
+        target: Target = 'update',
+    ):
         super().__init__(defaults)
         self._target: Target = target
         self._uses_grad = uses_grad
+        self._concat_params = concat_params
+        self._update_freq = update_freq
+        self._scale_first = scale_first
+        self._inner = inner
+
+    def update(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | None,
+        states: list[dict[str, Any]],
+        settings: Sequence[Mapping[str, Any]],
+    ) -> None:
+        """Updates this transform. By default does nothing - if logic is in `apply` method."""
 
     @abstractmethod
-    def transform(self, tensors: list[torch.Tensor], params: list[torch.Tensor], grads: list[torch.Tensor] | None, vars: Vars) -> Iterable[torch.Tensor]:
-        """applies the update rule to `target`."""
+    def apply(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | None,
+        states: list[dict[str, Any]],
+        settings: Sequence[Mapping[str, Any]],
+    ) -> Sequence[torch.Tensor]:
+        """Applies the update rule to `tensors`."""
+
+    @final
+    @torch.no_grad
+    def transform(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | None,
+        states: list[dict[str, Any]],
+        settings: Sequence[Mapping[str, Any]] | None,
+    ) -> list[torch.Tensor]:
+        """Applies this transform to an arbitrary sequence of tensors."""
+        un_tensors = tensors
+        un_params = params
+        un_grads = grads
+        if self._concat_params:
+            tensors = [torch.cat([t.ravel() for t in tensors])]
+            params = [torch.cat([p.ravel() for p in params])]
+            grads = [torch.cat([g.ravel() for g in grads])] if grads is not None else None
+
+        if settings is None:
+            settings = [self.defaults for _ in tensors]
+
+        step = self.global_state.get('__step', 0)
+        num = len(tensors)
+        states = states[:num]
+        settings = settings[:num]
+
+        update_freq = self._update_freq
+        scale_first = self._scale_first
+        scale_factor = 1
+
+        # scaling factor for 1st step
+        if scale_first and step == 0:
+            # initial step size guess from pytorch LBFGS
+            scale_factor = 1 / TensorList(tensors).abs().global_sum().clip(min=1)
+            scale_factor = scale_factor.clip(min=torch.finfo(tensors[0].dtype).eps)
+
+        # update transform
+        if step % update_freq == 0:
+            self.update(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+
+        # step with inner
+        if self._inner is not None:
+            tensors = apply_transform(self._inner, tensors=un_tensors, params=un_params, grads=un_grads)
+            if self._concat_params:
+                tensors = [torch.cat([t.ravel() for t in tensors])]
+
+        # apply transform
+        tensors = list(self.apply(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings))
+
+        # scale initial step, when preconditioner might not have been applied
+        if scale_first and step == 0:
+            torch._foreach_mul_(tensors, scale_factor)
+
+        self.global_state['__step'] = step + 1
+        if self._concat_params:
+            tensors = vec_to_tensors(vec=tensors[0], reference=un_tensors)
+        return tensors
+
 
-    def step(self, vars: Vars) -> Vars:
-        # vars may change, therefore current params and grads have to be extracted and passed explicitly
-        if self._uses_grad: vars.get_grad()
-        params=vars.params; grad = vars.grad
+    @torch.no_grad
+    def keyed_transform(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | None,
+    ):
+        """Applies this transform to `tensors`, `params` will be used as keys and need to always point to same tensor objects."""
+        if self._concat_params:
+            p = params[0]
+            states = [self.state[p]]
+            settings = [self.settings[p]]
+
+        else:
+            states = []
+            settings = []
+            for p in params:
+                states.append(self.state[p])
+                settings.append(self.settings[p])
+
+        return self.transform(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+
+    def step(self, var: Var) -> Var:
+        # var may change, therefore current params and grads have to be extracted and passed explicitly
+        if self._uses_grad: var.get_grad()
+        params=var.params
 
         # ---------------------------------- update ---------------------------------- #
         if self._target == 'update':
-            vars.update = list(self.transform(vars.get_update(), params, grad, vars))
-            return vars
+            update = var.get_update()
+            var.update = list(self.keyed_transform(update, params, var.grad, var.loss))
+            return var
 
         # ----------------------------------- grad ----------------------------------- #
         if self._target == 'grad':
-            vars.grad = list(self.transform(vars.get_grad(), params, grad, vars))
-            return vars
+            grad = var.get_grad()
+            var.grad = list(self.keyed_transform(grad, params, grad, var.loss))
+            return var
 
         # ------------------------------- params_direct ------------------------------ #
         if self._target == 'params_direct':
-            new_params = self.transform(vars.params, params, grad, vars)
-            for p, new_p in zip(vars.params, new_params): set_storage_(p, new_p)
-            return vars
+            new_params = self.keyed_transform(var.params, params, var.grad, var.loss)
+            for p, new_p in zip(var.params, new_params): set_storage_(p, new_p)
+            return var
 
         # ----------------------------- params_differnce ----------------------------- #
         if self._target == 'params_difference':
-            new_params = tuple(self.transform([p.clone() for p in vars.params], params, grad, vars))
-            vars.update = list(torch._foreach_sub(vars.params, new_params))
-            return vars
+            new_params = tuple(self.keyed_transform([p.clone() for p in var.params], params, var.grad, var.loss))
+            var.update = list(torch._foreach_sub(var.params, new_params))
+            return var
 
         # ----------------------------- update_difference ---------------------------- #
         if self._target == 'update_difference':
-            update = vars.get_update()
-            new_update = tuple(self.transform([u.clone() for u in update], params, grad, vars))
-            vars.update = list(torch._foreach_sub(update, new_update))
-            return vars
+            update = var.get_update()
+            new_update = tuple(self.keyed_transform([u.clone() for u in update], params, var.grad, var.loss))
+            var.update = list(torch._foreach_sub(update, new_update))
+            return var
 
         # ---------------------------------- closure --------------------------------- #
         if self._target == 'closure':
-            original_closure = vars.closure
+            original_closure = var.closure
             if original_closure is None: raise ValueError('Target = "closure", but closure is None')
 
-            params = vars.params
+            params = var.params
             def transformed_closure(backward=True):
                 if backward:
                     loss = original_closure()
                     current_grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
-                    transformed_grad = list(self.transform(current_grad, params, grad, vars))
+                    transformed_grad = list(self.keyed_transform(current_grad, params, var.grad, var.loss))
                     for p, g in zip(params, transformed_grad):
                         p.grad = g
 
@@ -84,14 +202,14 @@ class Transform(Module, ABC):
 
                 return loss
 
-            vars.closure = transformed_closure
-            return vars
+            var.closure = transformed_closure
+            return var
 
         # ---------------------------------- invalid --------------------------------- #
         raise ValueError(f'Invalid target: {self._target}')
 
 
-class TensorwiseTransform(Module, ABC):
+class TensorwiseTransform(Transform, ABC):
     """Base class for a parameter-wise transform.
 
     This is an abstract class, to use it, subclass it and override `transform`.
@@ -102,151 +220,94 @@ class TensorwiseTransform(Module, ABC):
             Set this to True if `transform` method uses the `grad` argument. This will ensure
             `grad` is always computed and can't be None. Otherwise set to False.
         target (Target, optional):
-            what to set on vars. Defaults to 'update'.
+            what to set on var. Defaults to 'update'.
     """
-    def __init__(self, defaults: dict[str,Any] | None, uses_grad: bool, target: Target = 'update'):
-        super().__init__(defaults)
-        self._target: Target = target
-        self._uses_grad: bool = uses_grad
+    def __init__(
+        self,
+        defaults: dict[str,Any] | None,
+        uses_grad: bool,
+        concat_params: bool = False,
+        update_freq: int = 1,
+        scale_first: bool = False,
+        inner: Chainable | None = None,
+        target: Target = 'update',
+    ):
+        super().__init__(
+            defaults=defaults,
+            uses_grad=uses_grad,
+            concat_params=concat_params,
+            update_freq=update_freq,
+            scale_first=scale_first,
+            inner=inner,
+            target=target,
+        )
+
+    def update_tensor(
+        self,
+        tensor: torch.Tensor,
+        param: torch.Tensor,
+        grad: torch.Tensor | None,
+        loss: torch.Tensor | None,
+        state: dict[str, Any],
+        settings: Mapping[str, Any],
+    ) -> None:
+        """Updates this transform. By default does nothing - if logic is in `apply` method."""
 
     @abstractmethod
-    def transform(
+    def apply_tensor(
         self,
         tensor: torch.Tensor,
         param: torch.Tensor,
         grad: torch.Tensor | None,
-        vars: Vars,
+        loss: torch.Tensor | None,
+        state: dict[str, Any],
+        settings: Mapping[str, Any],
     ) -> torch.Tensor:
-        """applies the update rule to `target`"""
-
-    def step(self, vars: Vars) -> Vars:
-        params = vars.params
-        if self._uses_grad and vars.grad is None: vars.get_grad()
-
-        # ---------------------------------- update ---------------------------------- #
-        if self._target == 'update':
-            update = vars.get_update()
-            grad = vars.grad if vars.grad is not None else [None] * len(params)
-            transformed_update = []
-
-            for p, g, u in zip(params, grad, update):
-                # settings = self.settings[p] # couldn't make typing work with this
-                #, self.transform(target=u, param=p, grad=g, vars=vars, **{k:settings[k] for k in self.defaults})
-                transformed_update.append(self.transform(tensor=u, param=p, grad=g, vars=vars))
-
-            vars.update = transformed_update
-            return vars
-
-        # ----------------------------------- grad ----------------------------------- #
-        if self._target == 'grad':
-            grad = vars.get_grad()
-            transformed_grad = []
-
-            for p, g in zip(params, grad):
-                transformed_grad.append(self.transform(tensor=g, param=p, grad=g, vars=vars))
-
-            vars.grad = transformed_grad
-            return vars
-
-        # ------------------------------- params_direct ------------------------------ #
-        if self._target == 'params_direct':
-            grad = vars.grad if vars.grad is not None else [None] * len(params)
-
-            for p, g in zip(params, grad):
-                set_storage_(p, self.transform(tensor=p, param=p, grad=g, vars=vars))
-
-            return vars
-
-        # ----------------------------- params_difference ---------------------------- #
-        if self._target == 'params_difference':
-            grad = vars.grad if vars.grad is not None else [None] * len(params)
-            transformed_params = []
-
-            for p, g in zip(params, grad):
-                transformed_params.append(
-                    self.transform(tensor=p.clone(), param=p, grad=g, vars=vars)
-                )
-
-            vars.update = list(torch._foreach_sub(params, transformed_params))
-            return vars
-
-        # ----------------------------- update_difference ---------------------------- #
-        if self._target == 'update_difference':
-            update = vars.get_update()
-            grad = vars.grad if vars.grad is not None else [None] * len(params)
-            transformed_update = []
-
-            for p, g, u in zip(params, grad, update):
-                transformed_update.append(
-                    self.transform(tensor=u.clone(), param=p, grad=g, vars=vars)
-                )
-
-            vars.update = list(torch._foreach_sub(update, transformed_update))
-            return vars
-
-        # ---------------------------------- closure --------------------------------- #
-        if self._target == 'closure':
-            original_closure = vars.closure
-            if original_closure is None: raise ValueError('Target = "closure", but closure is None')
-
-            params = vars.params
-            def transformed_closure(backward=True):
-                if backward:
-                    loss = original_closure()
-                    grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
-                    transformed_grad = []
-
-                    for p, g in zip(params, grad):
-                        transformed_grad.append(self.transform(tensor=g, param=p, grad=g, vars=vars))
-
-                    for p, g in zip(params, transformed_grad):
-                        p.grad = g
-
-                else:
-                    loss = original_closure(False)
-
-                return loss
-
-            vars.closure = transformed_closure
-            return vars
-
-        # ---------------------------------- invalid --------------------------------- #
-        raise ValueError(f'Invalid target: {self._target}')
-
-
-
-def apply(
+        """Applies the update rule to `tensor`."""
+
+    @final
+    def update(self, tensors, params, grads, loss, states, settings):
+        if grads is None: grads = [None]*len(tensors)
+        for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
+            self.update_tensor(t, p, g, loss, state, setting)
+
+    @final
+    def apply(self, tensors, params, grads, loss, states, settings):
+        applied = []
+        if grads is None: grads = [None]*len(tensors)
+        for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
+            applied.append(self.apply_tensor(t, p, g, loss, state, setting))
+        return applied
+
+def apply_transform(
     tfm: Chainable,
     tensors: list[torch.Tensor],
     params: list[torch.Tensor],
     grads: list[torch.Tensor] | None,
-    vars: Vars | None = None,
+    loss: torch.Tensor | None = None,
+    var: Var | None = None,
     current_step: int = 0,
 ):
-    if vars is None: vars = Vars(params=params, closure=None, model=None, current_step=current_step)
-    if isinstance(tfm, Transform):
-        if tfm._uses_grad and grads is None: grads = vars.get_grad()
-        return list(tfm.transform(tensors, params, grads, vars))
+    if var is None:
+        var = Var(params=params, closure=None, model=None, current_step=current_step)
+        var.loss = loss
 
-    if isinstance(tfm, TensorwiseTransform):
-        grads_list = grads
-        if grads_list is None:
-            if tfm._uses_grad: grads_list = vars.get_grad()
-            else: grads_list = [None] * len(tensors)
-        return [tfm.transform(t, p, g, vars) for t,p,g in zip(tensors,params,grads_list)]
+    if isinstance(tfm, Transform):
+        if tfm._uses_grad and grads is None: grads = var.get_grad()
+        return list(tfm.keyed_transform(tensors, params, grads, loss))
 
     if isinstance(tfm, Chain): tfm = tfm.get_children_sequence() # pyright: ignore[reportAssignmentType]
     if isinstance(tfm, Sequence):
         for module in tfm:
-            tensors = apply(module, tensors=tensors, params=params, grads=grads, vars=vars)
+            tensors = apply_transform(module, tensors=tensors, params=params, grads=grads, var=var)
         return tensors
 
     if isinstance(tfm, Module):
-        cvars = vars.clone(clone_update=False)
-        cvars.update = tensors
-        cvars = tfm.step(cvars)
-        vars.update_attrs_from_clone_(cvars)
-        assert cvars.update is not None
-        return cvars.update
+        cvar = var.clone(clone_update=False)
+        cvar.update = tensors
+        cvar = tfm.step(cvar)
+        var.update_attrs_from_clone_(cvar)
+        assert cvar.update is not None
+        return cvar.update
 
     raise TypeError(type(tfm))

torchzero/modules/__init__.py

@@ -11,3 +11,4 @@ from .smoothing import *
 from .weight_decay import *
 from .wrappers import *
 from .second_order import *
+from .higher_order import *

torchzero/modules/clipping/clipping.py

@@ -151,8 +151,8 @@ class ClipValue(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        value = self.get_settings('value', params=params)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        value = [s['value'] for s in settings]
         return TensorList(tensors).clip_([-v for v in value], value)
 
 class ClipNorm(Transform):
@@ -186,9 +186,9 @@ class ClipNorm(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        max_norm = self.get_settings('max_norm', params=params, cls=NumberList)
-        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(self.settings[params[0]])
+    def apply(self, tensors, params, grads, loss, states, settings):
+        max_norm = NumberList(s['max_norm'] for s in settings)
+        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(settings[0])
         _clip_norm_(
             tensors_ = TensorList(tensors),
             min = 0,
@@ -232,9 +232,9 @@ class Normalize(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        norm_value = self.get_settings('norm_value', params=params, cls=NumberList)
-        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(self.settings[params[0]])
+    def apply(self, tensors, params, grads, loss, states, settings):
+        norm_value = NumberList(s['norm_value'] for s in settings)
+        ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(settings[0])
 
         _clip_norm_(
             tensors_ = TensorList(tensors),
@@ -311,8 +311,8 @@ class Centralize(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        dim, min_size, inverse_dims = itemgetter('dim', 'min_size', 'inverse_dims')(self.settings[params[0]])
+    def apply(self, tensors, params, grads, loss, states, settings):
+        dim, min_size, inverse_dims = itemgetter('dim', 'min_size', 'inverse_dims')(settings[0])
 
         _centralize_(tensors_ = TensorList(tensors), dim=dim, inverse_dims=inverse_dims, min_size=min_size)
 

torchzero/modules/clipping/ema_clipping.py

@@ -4,8 +4,8 @@ from collections.abc import Iterable, Sequence
 
 import torch
 
-from ...core import Module, Target, Transform, apply, Chainable
-from ...utils import NumberList, TensorList, generic_eq
+from ...core import Module, Target, Transform, apply_transform, Chainable
+from ...utils import NumberList, TensorList, generic_eq, unpack_dicts, unpack_states
 
 class ClipNormByEMA(Transform):
     """Clips norm to be no larger than the norm of an exponential moving average of past updates.
@@ -34,13 +34,14 @@ class ClipNormByEMA(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        ord, tensorwise, ema_init, max_ema_growth = itemgetter('ord', 'tensorwise', 'ema_init', 'max_ema_growth')(self.settings[params[0]])
-
-        beta, eps = self.get_settings('beta', 'eps', params=params, cls=NumberList)
+    def apply(self, tensors, params, grads, loss, states, settings):
         tensors = TensorList(tensors)
+        ord, tensorwise, ema_init, max_ema_growth = itemgetter('ord', 'tensorwise', 'ema_init', 'max_ema_growth')(settings[0])
+
+        beta, eps = unpack_dicts(settings, 'beta', 'eps', cls=NumberList)
+
+        ema = unpack_states(states, tensors, 'ema', init = (torch.zeros_like if ema_init=='zeros' else tensors), cls=TensorList)
 
-        ema = self.get_state('ema', params=params, init = (torch.zeros_like if ema_init=='zeros' else tensors), cls=TensorList)
         ema.lerp_(tensors, 1-beta)
 
         if tensorwise:
@@ -48,7 +49,7 @@ class ClipNormByEMA(Transform):
 
             # clip ema norm growth
             if max_ema_growth is not None:
-                prev_ema_norm = self.get_state('prev_ema_norm', params=params, init=ema_norm, cls=TensorList)
+                prev_ema_norm = unpack_states(states, tensors, 'prev_ema_norm', init=ema_norm, cls=TensorList)
                 allowed_norm = (prev_ema_norm * max_ema_growth).clip(min=1e-6)
                 ema_denom = (ema_norm / allowed_norm).clip(min=1)
                 ema.div_(ema_denom)
@@ -119,17 +120,17 @@ class ClipValueByEMA(Transform):
             self.set_child('ema_tfm', ema_tfm)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        ema_init = itemgetter('ema_init')(self.settings[params[0]])
+    def apply(self, tensors, params, grads, loss, states, settings):
+        ema_init = itemgetter('ema_init')(settings[0])
 
-        beta = self.get_settings('beta', params=params, cls=NumberList)
+        beta = unpack_dicts(settings, 'beta', cls=NumberList)
         tensors = TensorList(tensors)
 
-        ema = self.get_state('ema', params=params, init = (torch.zeros_like if ema_init=='zeros' else lambda t: t.abs()), cls=TensorList)
+        ema = unpack_states(states, tensors, 'ema', init = (torch.zeros_like if ema_init=='zeros' else lambda t: t.abs()), cls=TensorList)
         ema.lerp_(tensors.abs(), 1-beta)
 
         if 'ema_tfm' in self.children:
-            ema = TensorList(apply(self.children['ema_tfm'], ema, params, vars.grad, vars))
+            ema = TensorList(apply_transform(self.children['ema_tfm'], ema, params, grads, loss))
 
         tensors.clip_(-ema, ema)
         return tensors

torchzero/modules/clipping/growth_clipping.py

@@ -19,7 +19,7 @@ class ClipValueGrowth(TensorwiseTransform):
             bounds the tracked multiplicative clipping decay to prevent collapse to 0.
             Next update is at most :code:`max(previous update * mul, max_decay)`.
             Defaults to 2.
-        target (Target, optional): what to set on vars.. Defaults to "update".
+        target (Target, optional): what to set on var.. Defaults to "update".
     """
     def __init__(
         self,
@@ -33,12 +33,10 @@ class ClipValueGrowth(TensorwiseTransform):
         super().__init__(defaults, uses_grad=False, target=target)
 
 
-    def transform(self, tensor, param, grad, vars):
-        add, mul, min_value, max_decay = itemgetter('add','mul','min_value','max_decay')(self.settings[param])
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+        add, mul, min_value, max_decay = itemgetter('add','mul','min_value','max_decay')(settings)
         add: float | None
 
-        state = self.state[param]
-
         if add is None and mul is None:
             return tensor
 
@@ -133,7 +131,7 @@ class ClipNormGrowth(Transform):
         ord (float, optional): norm order. Defaults to 2.
         parameterwise (bool, optional):
             if True, norms are calculated parameter-wise, otherwise treats all parameters as single vector. Defaults to True.
-        target (Target, optional): what to set on vars. Defaults to "update".
+        target (Target, optional): what to set on var. Defaults to "update".
     """
     def __init__(
         self,
@@ -150,35 +148,35 @@ class ClipNormGrowth(Transform):
 
 
 
-    def transform(self, tensors, params, grads, vars):
-        parameterwise = self.settings[params[0]]['parameterwise']
+    def apply(self, tensors, params, grads, loss, states, settings):
+        parameterwise = settings[0]['parameterwise']
         tensors = TensorList(tensors)
 
         if parameterwise:
             ts = tensors
-            stts = [self.state[p] for p in params]
-            stns = [self.settings[p] for p in params]
+            stts = states
+            stns = settings
 
         else:
             ts = [tensors.to_vec()]
             stts = [self.global_state]
-            stns = [self.settings[params[0]]]
+            stns = [settings[0]]
 
 
-        for t,state, settings in zip(ts, stts, stns):
+        for t, state, setting in zip(ts, stts, stns):
             if 'prev_norm' not in state:
-                state['prev_norm'] = torch.linalg.vector_norm(t, ord=settings['ord']) # pylint:disable=not-callable
+                state['prev_norm'] = torch.linalg.vector_norm(t, ord=setting['ord']) # pylint:disable=not-callable
                 state['prev_denom'] = 1
                 continue
 
             _,  state['prev_norm'], state['prev_denom'] = norm_growth_clip_(
                 tensor_ = t,
                 prev_norm = state['prev_norm'],
-                add = settings['add'],
-                mul = settings['mul'],
-                min_value = settings['min_value'],
-                max_decay = settings['max_decay'],
-                ord = settings['ord'],
+                add = setting['add'],
+                mul = setting['mul'],
+                min_value = setting['min_value'],
+                max_decay = setting['max_decay'],
+                ord = setting['ord'],
             )
 
         if not parameterwise: