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

torchzero/modules/ops/misc.py

@@ -5,8 +5,8 @@ from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, TensorwiseTransform, Target, Transform, Vars
-from ...utils import Distributions, NumberList, TensorList
+from ...core import Chainable, Module, Target, TensorwiseTransform, Transform, Var
+from ...utils import Distributions, NumberList, TensorList, unpack_dicts, unpack_states
 
 
 class Previous(TensorwiseTransform):
@@ -17,9 +17,8 @@ class Previous(TensorwiseTransform):
 
 
     @torch.no_grad
-    def transform(self, tensor, param, grad, vars):
-        n = self.settings[param]['n']
-        state = self.state[param]
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+        n = settings['n']
 
         if 'history' not in state:
             state['history'] = deque(maxlen=n+1)
@@ -35,10 +34,10 @@ class LastDifference(Transform):
         super().__init__({}, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        prev_target = self.get_state('prev_target', params=params) # initialized to 0
-        difference = torch._foreach_sub(tensors, prev_target)
-        for p, c in zip(prev_target, tensors): p.set_(c)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev_target') # initialized to 0
+        difference = torch._foreach_sub(tensors, prev)
+        for p, c in zip(prev, tensors): p.set_(c)
         return difference
 
 class LastGradDifference(Module):
@@ -47,13 +46,13 @@ class LastGradDifference(Module):
         super().__init__({})
 
     @torch.no_grad
-    def step(self, vars):
-        grad = vars.get_grad()
-        prev_grad = self.get_state('prev_grad', params=vars.params) # initialized to 0
+    def step(self, var):
+        grad = var.get_grad()
+        prev_grad = self.get_state(var.params, 'prev_grad') # initialized to 0
         difference = torch._foreach_sub(grad, prev_grad)
         for p, c in zip(prev_grad, grad): p.set_(c)
-        vars.update = list(difference)
-        return vars
+        var.update = list(difference)
+        return var
 
 
 class LastProduct(Transform):
@@ -62,10 +61,10 @@ class LastProduct(Transform):
         super().__init__({}, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        prev_target = self.get_state('prev_target', params=params, init=torch.ones_like) # initialized to 1 for prod
-        prod = torch._foreach_mul(tensors, prev_target)
-        for p, c in zip(prev_target, tensors): p.set_(c)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', init=torch.ones_like) # initialized to 1 for prod
+        prod = torch._foreach_mul(tensors, prev)
+        for p, c in zip(prev, tensors): p.set_(c)
         return prod
 
 class LastRatio(Transform):
@@ -75,12 +74,12 @@ class LastRatio(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        prev_target = self.get_state('prev_target', params=params, init = torch.ones_like) # initialized to ones
-        numerator = self.settings[params[0]]['numerator']
-        if numerator == 'cur': ratio = torch._foreach_div(tensors, prev_target)
-        else: ratio = torch._foreach_div(prev_target, tensors)
-        for p, c in zip(prev_target, tensors): p.set_(c)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', init = torch.ones_like) # initialized to ones
+        numerator = settings[0]['numerator']
+        if numerator == 'cur': ratio = torch._foreach_div(tensors, prev)
+        else: ratio = torch._foreach_div(prev, tensors)
+        for p, c in zip(prev, tensors): p.set_(c)
         return ratio
 
 class LastAbsoluteRatio(Transform):
@@ -90,17 +89,17 @@ class LastAbsoluteRatio(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        prev_target = self.get_state('prev_target', params=params, init = torch.ones_like) # initialized to 0
-        numerator = self.settings[params[0]]['numerator']
-        eps = self.get_settings('eps', params=params, cls = NumberList)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', init = torch.ones_like) # initialized to ones
+        numerator = settings[0]['numerator']
+        eps = NumberList(s['eps'] for s in settings)
 
         torch._foreach_abs_(tensors)
-        torch._foreach_clamp_min_(prev_target, eps)
+        torch._foreach_clamp_min_(prev, eps)
 
-        if numerator == 'cur': ratio = torch._foreach_div(tensors, prev_target)
-        else: ratio = torch._foreach_div(prev_target, tensors)
-        for p, c in zip(prev_target, tensors): p.set_(c)
+        if numerator == 'cur': ratio = torch._foreach_div(tensors, prev)
+        else: ratio = torch._foreach_div(prev, tensors)
+        for p, c in zip(prev, tensors): p.set_(c)
         return ratio
 
 class GradSign(Transform):
@@ -109,7 +108,7 @@ class GradSign(Transform):
         super().__init__({}, uses_grad=True, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         return [t.copysign_(g) for t,g in zip(tensors, grads)]
 
@@ -119,7 +118,7 @@ class UpdateSign(Transform):
         super().__init__({}, uses_grad=True, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         return [g.copysign(t) for t,g in zip(tensors, grads)] # no in-place
 
@@ -130,9 +129,9 @@ class GraftToGrad(Transform):
         super().__init__(defaults, uses_grad=True, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
-        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(self.settings[params[0]])
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(settings[0])
         return TensorList(tensors).graft_(grads, tensorwise=tensorwise, ord=ord, eps=eps)
 
 class GraftGradToUpdate(Transform):
@@ -142,9 +141,9 @@ class GraftGradToUpdate(Transform):
         super().__init__(defaults, uses_grad=True, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
-        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(self.settings[params[0]])
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(settings[0])
         return TensorList(grads).graft(tensors, tensorwise=tensorwise, ord=ord, eps=eps)
 
 
@@ -155,8 +154,8 @@ class GraftToParams(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(self.settings[params[0]])
+    def apply(self, tensors, params, grads, loss, states, settings):
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(settings[0])
         return TensorList(tensors).graft_(params, tensorwise=tensorwise, ord=ord, eps=eps)
 
 class Relative(Transform):
@@ -166,8 +165,8 @@ class Relative(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        mul = TensorList(params).abs().clamp_(self.get_settings('min_value', params=params))
+    def apply(self, tensors, params, grads, loss, states, settings):
+        mul = TensorList(params).abs().clamp_([s['min_value'] for s in settings])
         torch._foreach_mul_(tensors, mul)
         return tensors
 
@@ -178,94 +177,94 @@ class FillLoss(Module):
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        alpha = self.get_settings('alpha', params=vars.params)
-        loss = vars.get_loss(backward=self.settings[vars.params[0]]['backward'])
-        vars.update = [torch.full_like(p, loss*a) for p,a in zip(vars.params, alpha)]
-        return vars
+    def step(self, var):
+        alpha = self.get_settings(var.params, 'alpha')
+        loss = var.get_loss(backward=self.settings[var.params[0]]['backward'])
+        var.update = [torch.full_like(p, loss*a) for p,a in zip(var.params, alpha)]
+        return var
 
-class MulByLoss(Transform):
+class MulByLoss(Module):
     """multiplies update by loss times alpha"""
-    def __init__(self, alpha: float = 1, min_value:float = 1e-8, backward: bool = True, target: Target = 'update'):
+    def __init__(self, alpha: float = 1, min_value:float = 1e-8, backward: bool = True):
         defaults = dict(alpha=alpha, min_value=min_value, backward=backward)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars): #vars used for loss
-        alpha, min_value = self.get_settings('alpha', 'min_value', params=params)
-        loss = vars.get_loss(backward=self.settings[params[0]]['backward'])
+    def step(self, var):
+        alpha, min_value = self.get_settings(var.params, 'alpha', 'min_value')
+        loss = var.get_loss(backward=self.settings[var.params[0]]['backward'])
         mul = [max(loss*a, mv) for a,mv in zip(alpha, min_value)]
-        torch._foreach_mul_(tensors, mul)
-        return tensors
+        torch._foreach_mul_(var.update, mul)
+        return var
 
-class DivByLoss(Transform):
+class DivByLoss(Module):
     """divides update by loss times alpha"""
-    def __init__(self, alpha: float = 1, min_value:float = 1e-8, backward: bool = True, target: Target = 'update'):
+    def __init__(self, alpha: float = 1, min_value:float = 1e-8, backward: bool = True):
         defaults = dict(alpha=alpha, min_value=min_value, backward=backward)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars): #vars used for loss
-        alpha, min_value = self.get_settings('alpha', 'min_value', params=params)
-        loss = vars.get_loss(backward=self.settings[params[0]]['backward'])
+    def step(self, var):
+        alpha, min_value = self.get_settings(var.params, 'alpha', 'min_value')
+        loss = var.get_loss(backward=self.settings[var.params[0]]['backward'])
         mul = [max(loss*a, mv) for a,mv in zip(alpha, min_value)]
-        torch._foreach_div_(tensors, mul)
-        return tensors
+        torch._foreach_div_(var.update, mul)
+        return var
 
 
 
-def _sequential_step(self: Module, vars: Vars, sequential: bool):
-    params = vars.params
+def _sequential_step(self: Module, var: Var, sequential: bool):
+    params = var.params
     steps = self.settings[params[0]]['steps']
 
     if sequential: modules = self.get_children_sequence()
     else: modules = [self.children['module']] * steps
 
-    if vars.closure is None and len(modules) > 1: raise ValueError('Multistep and Sequential require closure')
+    if var.closure is None and len(modules) > 1: raise ValueError('Multistep and Sequential require closure')
 
     # store original params unless this is last module and can update params directly
-    params_before_steps = None if (vars.is_last and vars.last_module_lrs is None) else [p.clone() for p in params]
+    params_before_steps = None if (var.is_last and var.last_module_lrs is None) else [p.clone() for p in params]
 
-    # first step - pass vars as usual
-    vars = modules[0].step(vars)
-    new_vars = vars
+    # first step - pass var as usual
+    var = modules[0].step(var)
+    new_var = var
 
-    # subsequent steps - update parameters and create new vars
+    # subsequent steps - update parameters and create new var
     if len(modules) > 1:
         for m in modules[1:]:
 
             # update params
-            if (not new_vars.skip_update):
-                if new_vars.last_module_lrs is not None:
-                    torch._foreach_mul_(new_vars.get_update(), new_vars.last_module_lrs)
+            if (not new_var.skip_update):
+                if new_var.last_module_lrs is not None:
+                    torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
 
-                torch._foreach_sub_(params, new_vars.get_update())
+                torch._foreach_sub_(params, new_var.get_update())
 
-            # create new vars since we are at a new point, that means grad, update and loss will be None
-            new_vars = Vars(params=new_vars.params, closure=new_vars.closure,
-                            model=new_vars.model, current_step=new_vars.current_step + 1)
+            # create new var since we are at a new point, that means grad, update and loss will be None
+            new_var = Var(params=new_var.params, closure=new_var.closure,
+                            model=new_var.model, current_step=new_var.current_step + 1)
 
             # step
-            new_vars = m.step(new_vars)
+            new_var = m.step(new_var)
 
         # final parameter update
-        if (not new_vars.skip_update):
-            if new_vars.last_module_lrs is not None:
-                torch._foreach_mul_(new_vars.get_update(), new_vars.last_module_lrs)
+        if (not new_var.skip_update):
+            if new_var.last_module_lrs is not None:
+                torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
 
-            torch._foreach_sub_(params, new_vars.get_update())
+            torch._foreach_sub_(params, new_var.get_update())
 
-    # if last module, update is applied so return new vars
+    # if last module, update is applied so return new var
     if params_before_steps is None:
-        new_vars.stop = True
-        new_vars.skip_update = True
-        return new_vars
+        new_var.stop = True
+        new_var.skip_update = True
+        return new_var
 
     # otherwise use parameter difference as update
-    vars.update = list(torch._foreach_sub(params_before_steps, params))
+    var.update = list(torch._foreach_sub(params_before_steps, params))
     for p, bef in zip(params, params_before_steps):
         p.set_(bef) # pyright:ignore[reportArgumentType]
-    return vars
+    return var
 
 class Multistep(Module):
     def __init__(self, module: Chainable, steps: int):
@@ -274,8 +273,8 @@ class Multistep(Module):
         self.set_child('module', module)
 
     @torch.no_grad
-    def step(self, vars):
-        return _sequential_step(self, vars, sequential=False)
+    def step(self, var):
+        return _sequential_step(self, var, sequential=False)
 
 class Sequential(Module):
     def __init__(self, modules: Iterable[Chainable], steps: int):
@@ -284,8 +283,8 @@ class Sequential(Module):
         self.set_children_sequence(modules)
 
     @torch.no_grad
-    def step(self, vars):
-        return _sequential_step(self, vars, sequential=True)
+    def step(self, var):
+        return _sequential_step(self, var, sequential=True)
 
 
 class GradientAccumulation(Module):
@@ -297,22 +296,22 @@ class GradientAccumulation(Module):
 
 
     @torch.no_grad
-    def step(self, vars):
-        accumulator = self.get_state('accumulator', params=vars.params)
-        settings = self.settings[vars.params[0]]
+    def step(self, var):
+        accumulator = self.get_state(var.params, 'accumulator')
+        settings = self.settings[var.params[0]]
         n = settings['n']; mean = settings['mean']; stop = settings['stop']
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
         # add update to accumulator
-        torch._foreach_add_(accumulator, vars.get_update())
+        torch._foreach_add_(accumulator, var.get_update())
 
         # step with accumulated updates
         if step % n == 0:
             if mean:
                 torch._foreach_div_(accumulator, n)
 
-            vars.update = [a.clone() for a in accumulator]
-            vars = self.children['modules'].step(vars)
+            var.update = [a.clone() for a in accumulator]
+            var = self.children['modules'].step(var)
 
             # zero accumulator
             torch._foreach_zero_(accumulator)
@@ -320,10 +319,10 @@ class GradientAccumulation(Module):
         else:
             # prevent update
             if stop:
-                vars.stop=True
-                vars.skip_update=True
+                var.stop=True
+                var.skip_update=True
 
-        return vars
+        return var
 
 
 class Dropout(Transform):
@@ -332,10 +331,10 @@ class Dropout(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):
         tensors = TensorList(tensors)
-        p = self.get_settings('p', params=params, cls=NumberList)
-        graft = self.settings[params[0]]['graft']
+        p = NumberList(s['p'] for s in settings)
+        graft = settings[0]['graft']
 
         if graft:
             target_norm = tensors.global_vector_norm()
@@ -351,11 +350,11 @@ class WeightDropout(Module):
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        closure = vars.closure
+    def step(self, var):
+        closure = var.closure
         if closure is None: raise RuntimeError('WeightDropout requires closure')
-        params = TensorList(vars.params)
-        p = self.get_settings('p', params=params)
+        params = TensorList(var.params)
+        p = NumberList(self.settings[p]['p'] for p in params)
         mask = params.rademacher_like(p).add_(1).div_(2).as_bool()
 
         @torch.no_grad
@@ -369,8 +368,8 @@ class WeightDropout(Module):
             params.copy_(orig_params)
             return loss
 
-        vars.closure = dropout_closure
-        return vars
+        var.closure = dropout_closure
+        return var
 
 class NoiseSign(Transform):
     """uses random vector with update sign"""
@@ -379,8 +378,8 @@ class NoiseSign(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        alpha = self.get_settings('alpha', params=params)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        alpha = [s['alpha'] for s in settings]
         distribution = self.settings[params[0]]['distribution']
         return TensorList(tensors).sample_like(alpha, distribution).copysign_(tensors)
 
@@ -391,29 +390,29 @@ class NegateOnLossIncrease(Module):
         super().__init__(defaults=defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        closure = vars.closure
+    def step(self, var):
+        closure = var.closure
         if closure is None: raise RuntimeError('NegateOnLossIncrease requires closure')
-        backtrack = self.settings[vars.params[0]]['backtrack']
+        backtrack = self.settings[var.params[0]]['backtrack']
 
-        update = vars.get_update()
-        f_0 = vars.get_loss(backward=False)
+        update = var.get_update()
+        f_0 = var.get_loss(backward=False)
 
-        torch._foreach_sub_(vars.params, update)
+        torch._foreach_sub_(var.params, update)
         f_1 = closure(False)
 
         if f_1 <= f_0:
-            if vars.is_last and vars.last_module_lrs is None:
-                vars.stop = True
-                vars.skip_update = True
-                return vars
+            if var.is_last and var.last_module_lrs is None:
+                var.stop = True
+                var.skip_update = True
+                return var
 
-            torch._foreach_add_(vars.params, update)
-            return vars
+            torch._foreach_add_(var.params, update)
+            return var
 
-        torch._foreach_add_(vars.params, update)
+        torch._foreach_add_(var.params, update)
         if backtrack:
-            torch._foreach_neg_(vars.update)
+            torch._foreach_neg_(var.update)
         else:
-            torch._foreach_zero_(vars.update)
-        return vars
+            torch._foreach_zero_(var.update)
+        return var

torchzero/modules/ops/multi.py

@@ -7,7 +7,7 @@ from typing import Any
 
 import torch
 
-from ...core import Chainable, Module, Target, Vars, maybe_chain
+from ...core import Chainable, Module, Target, Var, maybe_chain
 from ...utils import TensorList, tensorlist
 
 
@@ -29,25 +29,25 @@ class MultiOperation(Module, ABC):
             raise ValueError('At least one operand must be a module')
 
     @abstractmethod
-    def transform(self, vars: Vars, **operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, **operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
         """applies the operation to operands"""
         raise NotImplementedError
 
     @torch.no_grad
-    def step(self, vars: Vars) -> Vars:
+    def step(self, var: Var) -> Var:
         # pass cloned update to all module operands
         processed_operands: dict[str, Any | list[torch.Tensor]] = self.operands.copy()
 
         for k,v in self.operands.items():
             if k in self.children:
                 v: Module
-                updated_vars = v.step(vars.clone(clone_update=True))
-                processed_operands[k] = updated_vars.get_update()
-                vars.update_attrs_from_clone_(updated_vars) # update loss, grad, etc if this module calculated them
+                updated_var = v.step(var.clone(clone_update=True))
+                processed_operands[k] = updated_var.get_update()
+                var.update_attrs_from_clone_(updated_var) # update loss, grad, etc if this module calculated them
 
-        transformed = self.transform(vars, **processed_operands)
-        vars.update = transformed
-        return vars
+        transformed = self.transform(var, **processed_operands)
+        var.update = transformed
+        return var
 
 
 
@@ -57,8 +57,8 @@ class SubModules(MultiOperation):
         super().__init__(defaults, input=input, other=other)
 
     @torch.no_grad
-    def transform(self, vars: Vars, input: float | list[torch.Tensor], other: float | list[torch.Tensor]) -> list[torch.Tensor]:
-        alpha = self.settings[vars.params[0]]['alpha']
+    def transform(self, var: Var, input: float | list[torch.Tensor], other: float | list[torch.Tensor]) -> list[torch.Tensor]:
+        alpha = self.settings[var.params[0]]['alpha']
 
         if isinstance(input, (int,float)):
             assert isinstance(other, list)
@@ -74,7 +74,7 @@ class DivModules(MultiOperation):
         super().__init__(defaults, input=input, other=other)
 
     @torch.no_grad
-    def transform(self, vars: Vars, input: float | list[torch.Tensor], other: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, input: float | list[torch.Tensor], other: float | list[torch.Tensor]) -> list[torch.Tensor]:
         if isinstance(input, (int,float)):
             assert isinstance(other, list)
             return input / TensorList(other)
@@ -88,7 +88,7 @@ class PowModules(MultiOperation):
         super().__init__(defaults, input=input, exponent=exponent)
 
     @torch.no_grad
-    def transform(self, vars: Vars, input: float | list[torch.Tensor], exponent: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, input: float | list[torch.Tensor], exponent: float | list[torch.Tensor]) -> list[torch.Tensor]:
         if isinstance(input, (int,float)):
             assert isinstance(exponent, list)
             return input ** TensorList(exponent)
@@ -102,8 +102,8 @@ class LerpModules(MultiOperation):
         super().__init__(defaults, input=input, end=end)
 
     @torch.no_grad
-    def transform(self, vars: Vars, input: list[torch.Tensor], end: list[torch.Tensor]) -> list[torch.Tensor]:
-        torch._foreach_lerp_(input, end, weight=self.settings[vars.params[0]]['weight'])
+    def transform(self, var: Var, input: list[torch.Tensor], end: list[torch.Tensor]) -> list[torch.Tensor]:
+        torch._foreach_lerp_(input, end, weight=self.settings[var.params[0]]['weight'])
         return input
 
 class ClipModules(MultiOperation):
@@ -112,7 +112,7 @@ class ClipModules(MultiOperation):
         super().__init__(defaults, input=input, min=min, max=max)
 
     @torch.no_grad
-    def transform(self, vars: Vars, input: list[torch.Tensor], min: float | list[torch.Tensor], max: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, var: Var, input: list[torch.Tensor], min: float | list[torch.Tensor], max: float | list[torch.Tensor]) -> list[torch.Tensor]:
         return TensorList(input).clamp_(min=min, max=max)
 
 
@@ -122,8 +122,8 @@ class GraftModules(MultiOperation):
         super().__init__(defaults, direction=direction, magnitude=magnitude)
 
     @torch.no_grad
-    def transform(self, vars, magnitude: list[torch.Tensor], direction:list[torch.Tensor]):
-        tensorwise, ord, eps, strength = itemgetter('tensorwise','ord','eps', 'strength')(self.settings[vars.params[0]])
+    def transform(self, var, magnitude: list[torch.Tensor], direction:list[torch.Tensor]):
+        tensorwise, ord, eps, strength = itemgetter('tensorwise','ord','eps', 'strength')(self.settings[var.params[0]])
         return TensorList(direction).graft_(magnitude, tensorwise=tensorwise, ord=ord, eps=eps, strength=strength)
 
 
@@ -132,6 +132,6 @@ class Where(MultiOperation):
         super().__init__({}, condition=condition, input=input, other=other)
 
     @torch.no_grad
-    def transform(self, vars, condition: list[torch.Tensor], input: list[torch.Tensor] | float, other: list[torch.Tensor] | float):
+    def transform(self, var, condition: list[torch.Tensor], input: list[torch.Tensor] | float, other: list[torch.Tensor] | float):
         return tensorlist.where(TensorList(condition).as_bool(), input, other)