torchzero 0.3.15__py3-none-any.whl → 0.4.1__py3-none-any.whl

torchzero/modules/ops/multi.py

@@ -7,8 +7,8 @@ from typing import Any, Literal
 
 import torch
 
-from ...core import Chainable, Module, Target, Var, maybe_chain
-from ...utils import TensorList, tensorlist, Metrics
+from ...core import Chainable, Module,  Objective
+from ...utils import TensorList, Metrics
 
 
 class MultiOperationBase(Module, ABC):
@@ -29,36 +29,39 @@ class MultiOperationBase(Module, ABC):
             raise ValueError('At least one operand must be a module')
 
     @abstractmethod
-    def transform(self, var: Var, **operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, **operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
         """applies the operation to operands"""
         raise NotImplementedError
 
+    def update(self, objective): raise RuntimeError
+    def apply(self, objective): raise RuntimeError
+
     @torch.no_grad
-    def step(self, var: Var) -> Var:
+    def step(self, objective: Objective) -> Objective:
         # 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_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
+                updated_obj = v.step(objective.clone(clone_updates=True))
+                processed_operands[k] = updated_obj.get_updates()
+                objective.update_attrs_from_clone_(updated_obj) # update loss, grad, etc if this module calculated them
 
-        transformed = self.transform(var, **processed_operands)
-        var.update = transformed
-        return var
+        transformed = self.transform(objective, **processed_operands)
+        objective.updates = transformed
+        return objective
 
 
 
 class SubModules(MultiOperationBase):
-    """Calculates :code:`input - other`. :code:`input` and :code:`other` can be numbers or modules."""
+    """Calculates ``input - other``. ``input`` and ``other`` can be numbers or modules."""
     def __init__(self, input: Chainable | float, other: Chainable | float, alpha: float = 1):
         defaults = dict(alpha=alpha)
         super().__init__(defaults, input=input, other=other)
 
     @torch.no_grad
-    def transform(self, var: Var, input: float | list[torch.Tensor], other: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, input: float | list[torch.Tensor], other: float | list[torch.Tensor]) -> list[torch.Tensor]:
         alpha = self.defaults['alpha']
 
         if isinstance(input, (int,float)):
@@ -70,14 +73,14 @@ class SubModules(MultiOperationBase):
         return input
 
 class DivModules(MultiOperationBase):
-    """Calculates :code:`input / other`. :code:`input` and :code:`other` can be numbers or modules."""
+    """Calculates ``input / other``. ``input`` and ``other`` can be numbers or modules."""
     def __init__(self, input: Chainable | float, other: Chainable | float, other_first:bool=False):
         defaults = {}
         if other_first: super().__init__(defaults, other=other, input=input)
         else: super().__init__(defaults, input=input, other=other)
 
     @torch.no_grad
-    def transform(self, var: Var, input: float | list[torch.Tensor], other: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, 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)
@@ -87,13 +90,13 @@ class DivModules(MultiOperationBase):
 
 
 class PowModules(MultiOperationBase):
-    """Calculates :code:`input ** exponent`. :code:`input` and :code:`other` can be numbers or modules."""
+    """Calculates ``input ** exponent``. ``input`` and ``other`` can be numbers or modules."""
     def __init__(self, input: Chainable | float, exponent: Chainable | float):
         defaults = {}
         super().__init__(defaults, input=input, exponent=exponent)
 
     @torch.no_grad
-    def transform(self, var: Var, input: float | list[torch.Tensor], exponent: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, 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,32 +105,32 @@ class PowModules(MultiOperationBase):
         return input
 
 class LerpModules(MultiOperationBase):
-    """Does a linear interpolation of :code:`input(tensors)` and :code:`end(tensors)` based on a scalar :code:`weight`.
+    """Does a linear interpolation of ``input(tensors)`` and ``end(tensors)`` based on a scalar ``weight``.
 
-    The output is given by :code:`output = input(tensors) + weight * (end(tensors) - input(tensors))`
+    The output is given by ``output = input(tensors) + weight * (end(tensors) - input(tensors))``
     """
     def __init__(self, input: Chainable, end: Chainable, weight: float):
         defaults = dict(weight=weight)
         super().__init__(defaults, input=input, end=end)
 
     @torch.no_grad
-    def transform(self, var: Var, input: list[torch.Tensor], end: list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, input: list[torch.Tensor], end: list[torch.Tensor]) -> list[torch.Tensor]:
         torch._foreach_lerp_(input, end, weight=self.defaults['weight'])
         return input
 
 class ClipModules(MultiOperationBase):
-    """Calculates :code:`input(tensors).clip(min, max)`. :code:`min` and :code:`max` can be numbers or modules."""
+    """Calculates ``input(tensors).clip(min, max)``. ``min`` and ``max`` can be numbers or modules."""
     def __init__(self, input: Chainable, min: float | Chainable | None = None, max: float | Chainable | None = None):
         defaults = {}
         super().__init__(defaults, input=input, min=min, max=max)
 
     @torch.no_grad
-    def transform(self, var: Var, input: list[torch.Tensor], min: float | list[torch.Tensor], max: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, 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)
 
 
-class GraftModules(MultiOperationBase):
-    """Outputs :code:`direction` output rescaled to have the same norm as :code:`magnitude` output.
+class Graft(MultiOperationBase):
+    """Outputs ``direction`` output rescaled to have the same norm as ``magnitude`` output.
 
     Args:
         direction (Chainable): module to use the direction from
@@ -137,40 +140,40 @@ class GraftModules(MultiOperationBase):
         eps (float, optional): clips denominator to be no less than this value. Defaults to 1e-6.
         strength (float, optional): strength of grafting. Defaults to 1.
 
-    Example:
-        Shampoo grafted to Adam
-
-        .. code-block:: python
+    ### Example:
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.GraftModules(
-                    direction = tz.m.Shampoo(),
-                    magnitude = tz.m.Adam(),
-                ),
-                tz.m.LR(1e-3)
-            )
+    Shampoo grafted to Adam
+    ```python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.GraftModules(
+            direction = tz.m.Shampoo(),
+            magnitude = tz.m.Adam(),
+        ),
+        tz.m.LR(1e-3)
+    )
+    ```
 
     Reference:
-        Agarwal, N., Anil, R., Hazan, E., Koren, T., & Zhang, C. (2020). Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11803. https://arxiv.org/pdf/2002.11803
+        [Agarwal, N., Anil, R., Hazan, E., Koren, T., & Zhang, C. (2020). Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11803.](https://arxiv.org/pdf/2002.11803)
     """
     def __init__(self, direction: Chainable, magnitude: Chainable, tensorwise:bool=True, ord:Metrics=2, eps:float = 1e-6, strength:float=1):
         defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps, strength=strength)
         super().__init__(defaults, direction=direction, magnitude=magnitude)
 
     @torch.no_grad
-    def transform(self, var, magnitude: list[torch.Tensor], direction:list[torch.Tensor]):
+    def transform(self, objective, magnitude: list[torch.Tensor], direction:list[torch.Tensor]):
         tensorwise, ord, eps, strength = itemgetter('tensorwise','ord','eps', 'strength')(self.defaults)
         return TensorList(direction).graft_(magnitude, tensorwise=tensorwise, ord=ord, eps=eps, strength=strength)
 
 class MultiplyByModuleNorm(MultiOperationBase):
-    """Outputs :code:`input` multiplied by norm of the :code:`norm` output."""
+    """Outputs ``input`` multiplied by norm of the ``norm`` output."""
     def __init__(self, input: Chainable, norm: Chainable, tensorwise:bool=True, ord:Metrics=2):
         defaults = dict(tensorwise=tensorwise, ord=ord)
         super().__init__(defaults, input=input, norm=norm)
 
     @torch.no_grad
-    def transform(self, var, input: list[torch.Tensor], norm:list[torch.Tensor]):
+    def transform(self, objective, input: list[torch.Tensor], norm:list[torch.Tensor]):
         tensorwise, ord = itemgetter('tensorwise','ord')(self.defaults)
         if tensorwise:
             n = TensorList(norm).metric(ord)
@@ -181,13 +184,13 @@ class MultiplyByModuleNorm(MultiOperationBase):
         return input
 
 class DivideByModuleNorm(MultiOperationBase):
-    """Outputs :code:`input` divided by norm of the :code:`norm` output."""
+    """Outputs ``input`` divided by norm of the ``norm`` output."""
     def __init__(self, input: Chainable, norm: Chainable, tensorwise:bool=True, ord:Metrics=2):
         defaults = dict(tensorwise=tensorwise, ord=ord)
         super().__init__(defaults, input=input, norm=norm)
 
     @torch.no_grad
-    def transform(self, var, input: list[torch.Tensor], norm:list[torch.Tensor]):
+    def transform(self, objective, input: list[torch.Tensor], norm:list[torch.Tensor]):
         tensorwise, ord = itemgetter('tensorwise','ord')(self.defaults)
         if tensorwise:
             n = TensorList(norm).metric(ord)

torchzero/modules/ops/reduce.py

@@ -5,7 +5,7 @@ from typing import Any, cast
 
 import torch
 
-from ...core import Chainable, Module, Target, Var, maybe_chain
+from ...core import Chainable, Module,  Objective, maybe_chain
 
 
 class ReduceOperationBase(Module, ABC):
@@ -26,34 +26,37 @@ class ReduceOperationBase(Module, ABC):
             raise ValueError('At least one operand must be a module')
 
     @abstractmethod
-    def transform(self, var: Var, *operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, *operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
         """applies the operation to operands"""
         raise NotImplementedError
 
+    def update(self, objective): raise RuntimeError
+    def apply(self, objective): raise RuntimeError
+
     @torch.no_grad
-    def step(self, var: Var) -> Var:
+    def step(self, objective: Objective) -> Objective:
         # pass cloned update to all module operands
         processed_operands: list[Any | list[torch.Tensor]] = self.operands.copy()
 
         for i, v in enumerate(self.operands):
             if f'operand_{i}' in self.children:
                 v: Module
-                updated_var = v.step(var.clone(clone_update=True))
-                processed_operands[i] = updated_var.get_update()
-                var.update_attrs_from_clone_(updated_var) # update loss, grad, etc if this module calculated them
+                updated_obj = v.step(objective.clone(clone_updates=True))
+                processed_operands[i] = updated_obj.get_updates()
+                objective.update_attrs_from_clone_(updated_obj) # update loss, grad, etc if this module calculated them
 
-        transformed = self.transform(var, *processed_operands)
-        var.update = transformed
-        return var
+        transformed = self.transform(objective, *processed_operands)
+        objective.updates = transformed
+        return objective
 
 class Sum(ReduceOperationBase):
-    """Outputs sum of :code:`inputs` that can be modules or numbers."""
+    """Outputs sum of ``inputs`` that can be modules or numbers."""
     USE_MEAN = False
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         sum = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:
@@ -64,14 +67,14 @@ class Sum(ReduceOperationBase):
         return sum
 
 class Mean(Sum):
-    """Outputs a mean of :code:`inputs` that can be modules or numbers."""
+    """Outputs a mean of ``inputs`` that can be modules or numbers."""
     USE_MEAN = True
 
 
 class WeightedSum(ReduceOperationBase):
+    """Outputs a weighted sum of ``inputs`` that can be modules or numbers."""
     USE_MEAN = False
     def __init__(self, *inputs: Chainable | float, weights: Iterable[float]):
-        """Outputs a weighted sum of :code:`inputs` that can be modules or numbers."""
         weights = list(weights)
         if len(inputs) != len(weights):
             raise ValueError(f'Number of inputs {len(inputs)} must match number of weights {len(weights)}')
@@ -79,7 +82,7 @@ class WeightedSum(ReduceOperationBase):
         super().__init__(defaults=defaults, *inputs)
 
     @torch.no_grad
-    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         weights = self.defaults['weights']
         sum = cast(list, sorted_inputs[0])
@@ -94,16 +97,16 @@ class WeightedSum(ReduceOperationBase):
 
 
 class WeightedMean(WeightedSum):
-    """Outputs weighted mean of :code:`inputs` that can be modules or numbers."""
+    """Outputs weighted mean of ``inputs`` that can be modules or numbers."""
     USE_MEAN = True
 
 class Median(ReduceOperationBase):
-    """Outputs median of :code:`inputs` that can be modules or numbers."""
+    """Outputs median of ``inputs`` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         res = []
         lists = [i for i in inputs if isinstance(i, list)]
         floats = [i for i in inputs if isinstance(i, (int,float))]
@@ -112,12 +115,12 @@ class Median(ReduceOperationBase):
         return res
 
 class Prod(ReduceOperationBase):
-    """Outputs product of :code:`inputs` that can be modules or numbers."""
+    """Outputs product of ``inputs`` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         prod = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:
@@ -127,12 +130,12 @@ class Prod(ReduceOperationBase):
         return prod
 
 class MaximumModules(ReduceOperationBase):
-    """Outputs elementwise maximum of :code:`inputs` that can be modules or numbers."""
+    """Outputs elementwise maximum of ``inputs`` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         maximum = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:
@@ -142,12 +145,12 @@ class MaximumModules(ReduceOperationBase):
         return maximum
 
 class MinimumModules(ReduceOperationBase):
-    """Outputs elementwise minimum of :code:`inputs` that can be modules or numbers."""
+    """Outputs elementwise minimum of ``inputs`` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
     @torch.no_grad
-    def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+    def transform(self, objective: Objective, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
         minimum = cast(list, sorted_inputs[0])
         if len(sorted_inputs) > 1:

torchzero/modules/ops/unary.py

@@ -2,102 +2,102 @@ from collections import deque
 
 import torch
 
-from ...core import TensorwiseTransform, Target, Transform
+from ...core import TensorTransform
 from ...utils import TensorList, unpack_dicts,unpack_states
 
-class UnaryLambda(Transform):
-    """Applies :code:`fn` to input tensors.
+class UnaryLambda(TensorTransform):
+    """Applies ``fn`` to input tensors.
 
-    :code:`fn` must accept and return a list of tensors.
+    ``fn`` must accept and return a list of tensors.
     """
-    def __init__(self, fn, target: "Target" = 'update'):
+    def __init__(self, fn):
         defaults = dict(fn=fn)
-        super().__init__(defaults=defaults, uses_grad=False, target=target)
+        super().__init__(defaults=defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         return settings[0]['fn'](tensors)
 
-class UnaryParameterwiseLambda(TensorwiseTransform):
-    """Applies :code:`fn` to each input tensor.
+class UnaryParameterwiseLambda(TensorTransform):
+    """Applies ``fn`` to each input tensor.
 
-    :code:`fn` must accept and return a tensor.
+    ``fn`` must accept and return a tensor.
     """
-    def __init__(self, fn, target: "Target" = 'update'):
+    def __init__(self, fn):
         defaults = dict(fn=fn)
-        super().__init__(uses_grad=False, defaults=defaults, target=target)
+        super().__init__(defaults=defaults)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
         return setting['fn'](tensor)
 
-class CustomUnaryOperation(Transform):
-    """Applies :code:`getattr(tensor, name)` to each tensor
+class CustomUnaryOperation(TensorTransform):
+    """Applies ``getattr(tensor, name)`` to each tensor
     """
-    def __init__(self, name: str, target: "Target" = 'update'):
+    def __init__(self, name: str):
         defaults = dict(name=name)
-        super().__init__(defaults=defaults, uses_grad=False, target=target)
+        super().__init__(defaults=defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         return getattr(tensors, settings[0]['name'])()
 
 
-class Abs(Transform):
-    """Returns :code:`abs(input)`"""
-    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+class Abs(TensorTransform):
+    """Returns ``abs(input)``"""
+    def __init__(self): super().__init__()
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_abs_(tensors)
         return tensors
 
-class Sign(Transform):
-    """Returns :code:`sign(input)`"""
-    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+class Sign(TensorTransform):
+    """Returns ``sign(input)``"""
+    def __init__(self): super().__init__()
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_sign_(tensors)
         return tensors
 
-class Exp(Transform):
-    """Returns :code:`exp(input)`"""
-    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+class Exp(TensorTransform):
+    """Returns ``exp(input)``"""
+    def __init__(self): super().__init__()
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_exp_(tensors)
         return tensors
 
-class Sqrt(Transform):
-    """Returns :code:`sqrt(input)`"""
-    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+class Sqrt(TensorTransform):
+    """Returns ``sqrt(input)``"""
+    def __init__(self): super().__init__()
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_sqrt_(tensors)
         return tensors
 
-class Reciprocal(Transform):
-    """Returns :code:`1 / input`"""
-    def __init__(self, eps = 0, target: "Target" = 'update'):
+class Reciprocal(TensorTransform):
+    """Returns ``1 / input``"""
+    def __init__(self, eps = 0):
         defaults = dict(eps = eps)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         eps = [s['eps'] for s in settings]
         if any(e != 0 for e in eps): torch._foreach_add_(tensors, eps)
         torch._foreach_reciprocal_(tensors)
         return tensors
 
-class Negate(Transform):
-    """Returns :code:`- input`"""
-    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+class Negate(TensorTransform):
+    """Returns ``- input``"""
+    def __init__(self): super().__init__()
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         torch._foreach_neg_(tensors)
         return tensors
 
 
-class NanToNum(Transform):
-    """Convert `nan`, `inf` and `-inf` to numbers.
+class NanToNum(TensorTransform):
+    """Convert ``nan``, ``inf`` and `-`inf`` to numbers.
 
     Args:
         nan (optional): the value to replace NaNs with. Default is zero.
@@ -108,23 +108,23 @@ class NanToNum(Transform):
             If None, negative infinity values are replaced with the lowest finite value
             representable by input's dtype. Default is None.
     """
-    def __init__(self, nan=None, posinf=None, neginf=None, target: "Target" = 'update'):
+    def __init__(self, nan=None, posinf=None, neginf=None):
         defaults = dict(nan=nan, posinf=posinf, neginf=neginf)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         nan, posinf, neginf = unpack_dicts(settings, 'nan', 'posinf', 'neginf')
         return [t.nan_to_num_(nan_i, posinf_i, neginf_i) for t, nan_i, posinf_i, neginf_i in zip(tensors, nan, posinf, neginf)]
 
-class Rescale(Transform):
-    """Rescales input to :code`(min, max)` range"""
-    def __init__(self, min: float, max: float, tensorwise: bool = False, eps:float=1e-8, target: "Target" = 'update'):
+class Rescale(TensorTransform):
+    """Rescales input to ``(min, max)`` range"""
+    def __init__(self, min: float, max: float, tensorwise: bool = False, eps:float=1e-8):
         defaults = dict(min=min, max=max, eps=eps, tensorwise=tensorwise)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         min, max = unpack_dicts(settings, 'min','max')
         tensorwise = settings[0]['tensorwise']
         dim = None if tensorwise else 'global'