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

torchzero/modules/momentum/averaging.py

@@ -1,27 +1,27 @@
 """Modules that perform averaging over a history of past updates."""
 from collections import deque
 from collections.abc import Sequence
-from typing import Any, Literal, cast
+from typing import Any
 
 import torch
 
-from ...core import TensorwiseTransform, Target
+from ...core import TensorTransform
 from ...utils import tolist
 
 
-class Averaging(TensorwiseTransform):
+class Averaging(TensorTransform):
     """Average of past ``history_size`` updates.
 
     Args:
         history_size (int): Number of past updates to average
         target (Target, optional): target. Defaults to 'update'.
     """
-    def __init__(self, history_size: int, target: Target = 'update'):
+    def __init__(self, history_size: int):
         defaults = dict(history_size=history_size)
-        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):
         history_size = setting['history_size']
         if 'history' not in state:
             state['history'] = deque(maxlen=history_size)
@@ -34,19 +34,19 @@ class Averaging(TensorwiseTransform):
 
         return average / len(history)
 
-class WeightedAveraging(TensorwiseTransform):
+class WeightedAveraging(TensorTransform):
     """Weighted average of past ``len(weights)`` updates.
 
     Args:
         weights (Sequence[float]): a sequence of weights from oldest to newest.
         target (Target, optional): target. Defaults to 'update'.
     """
-    def __init__(self, weights: Sequence[float] | torch.Tensor | Any, target: Target = 'update'):
+    def __init__(self, weights: Sequence[float] | torch.Tensor | Any):
         defaults = dict(weights = tolist(weights))
-        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):
         weights = setting['weights']
 
         if 'history' not in state:
@@ -68,19 +68,19 @@ class WeightedAveraging(TensorwiseTransform):
         return average
 
 
-class MedianAveraging(TensorwiseTransform):
+class MedianAveraging(TensorTransform):
     """Median of past ``history_size`` updates.
 
     Args:
         history_size (int): Number of past updates to average
         target (Target, optional): target. Defaults to 'update'.
     """
-    def __init__(self, history_size: int, target: Target = 'update'):
+    def __init__(self, history_size: int,):
         defaults = dict(history_size = history_size)
-        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):
         history_size = setting['history_size']
 
         if 'history' not in state:

torchzero/modules/momentum/cautious.py

@@ -5,7 +5,7 @@ from typing import Literal
 
 import torch
 
-from ...core import Target, Transform, Module, Chainable
+from ...core import  TensorTransform, Module, Chainable
 from ...utils import NumberList, TensorList, unpack_dicts
 
 
@@ -36,7 +36,7 @@ def cautious_(
     tensors_ -= tensors_.mul(2).mul_(mask.logical_not_())
     return tensors_
 
-class Cautious(Transform):
+class Cautious(TensorTransform):
     """Negates update for parameters where update and gradient sign is inconsistent.
     Optionally normalizes the update by the number of parameters that are not masked.
     This is meant to be used after any momentum-based modules.
@@ -57,7 +57,7 @@ class Cautious(Transform):
     Cautious Adam
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         bench.parameters(),
         tz.m.Adam(),
         tz.m.Cautious(),
@@ -79,12 +79,12 @@ class Cautious(Transform):
         super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(settings[0])
         return cautious_(TensorList(tensors), TensorList(grads), normalize=normalize, eps=eps, mode=mode)
 
-class UpdateGradientSignConsistency(Transform):
+class UpdateGradientSignConsistency(TensorTransform):
     """Compares update and gradient signs. Output will have 1s where signs match, and 0s where they don't.
 
     Args:
@@ -98,7 +98,7 @@ class UpdateGradientSignConsistency(Transform):
         super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         normalize, eps = itemgetter('normalize', 'eps')(settings[0])
 
@@ -108,7 +108,7 @@ 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.
+    """Negaties update on :code:`main` module where it's sign doesn't match with output of ``compare`` module.
 
     Args:
         main (Chainable): main module or sequence of modules whose update will be cautioned.
@@ -137,29 +137,32 @@ class IntermoduleCautious(Module):
         self.set_child('main', main)
         self.set_child('compare', compare)
 
+    def update(self, objective): raise RuntimeError
+    def apply(self, objective): raise RuntimeError
+
     @torch.no_grad
-    def step(self, var):
+    def step(self, objective):
         main = self.children['main']
         compare = self.children['compare']
 
-        main_var = main.step(var.clone(clone_update=True))
-        var.update_attrs_from_clone_(main_var)
+        main_var = main.step(objective.clone(clone_updates=True))
+        objective.update_attrs_from_clone_(main_var)
 
-        compare_var = compare.step(var.clone(clone_update=True))
-        var.update_attrs_from_clone_(compare_var)
+        compare_var = compare.step(objective.clone(clone_updates=True))
+        objective.update_attrs_from_clone_(compare_var)
 
         mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(self.defaults)
-        var.update = cautious_(
-            TensorList(main_var.get_update()),
-            TensorList(compare_var.get_update()),
+        objective.updates = cautious_(
+            TensorList(main_var.get_updates()),
+            TensorList(compare_var.get_updates()),
             normalize=normalize,
             mode=mode,
             eps=eps,
         )
 
-        return var
+        return objective
 
-class ScaleByGradCosineSimilarity(Transform):
+class ScaleByGradCosineSimilarity(TensorTransform):
     """Multiplies the update by cosine similarity with gradient.
     If cosine similarity is negative, naturally the update will be negated as well.
 
@@ -170,7 +173,7 @@ class ScaleByGradCosineSimilarity(Transform):
 
     Scaled Adam
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         bench.parameters(),
         tz.m.Adam(),
         tz.m.ScaleByGradCosineSimilarity(),
@@ -186,7 +189,7 @@ class ScaleByGradCosineSimilarity(Transform):
         super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         eps = settings[0]['eps']
         tensors = TensorList(tensors)
@@ -196,8 +199,8 @@ 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.
+    """Scales the output of ``main`` module by it's cosine similarity to the output
+    of ``compare`` module.
 
     Args:
         main (Chainable): main module or sequence of modules whose update will be scaled.
@@ -208,7 +211,7 @@ class ScaleModulesByCosineSimilarity(Module):
 
     Adam scaled by similarity to RMSprop
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         bench.parameters(),
         tz.m.ScaleModulesByCosineSimilarity(
             main = tz.m.Adam(),
@@ -230,22 +233,25 @@ class ScaleModulesByCosineSimilarity(Module):
         self.set_child('main', main)
         self.set_child('compare', compare)
 
+    def update(self, objective): raise RuntimeError
+    def apply(self, objective): raise RuntimeError
+
     @torch.no_grad
-    def step(self, var):
+    def step(self, objective):
         main = self.children['main']
         compare = self.children['compare']
 
-        main_var = main.step(var.clone(clone_update=True))
-        var.update_attrs_from_clone_(main_var)
+        main_var = main.step(objective.clone(clone_updates=True))
+        objective.update_attrs_from_clone_(main_var)
 
-        compare_var = compare.step(var.clone(clone_update=True))
-        var.update_attrs_from_clone_(compare_var)
+        compare_var = compare.step(objective.clone(clone_updates=True))
+        objective.update_attrs_from_clone_(compare_var)
 
-        m = TensorList(main_var.get_update())
-        c = TensorList(compare_var.get_update())
+        m = TensorList(main_var.get_updates())
+        c = TensorList(compare_var.get_updates())
         eps = self.defaults['eps']
 
         cos_sim = m.dot(c) / (m.global_vector_norm() * c.global_vector_norm()).clip(min=eps)
 
-        var.update = m.mul_(cos_sim)
-        return var
+        objective.updates = m.mul_(cos_sim)
+        return objective

torchzero/modules/momentum/momentum.py

@@ -4,12 +4,12 @@ from typing import Literal
 
 import torch
 
-from ...core import Target, Transform
+from ...core import  TensorTransform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
-from ..functional import debias, ema_
+from ..opt_utils import debias, ema_
 
 
-class EMA(Transform):
+class EMA(TensorTransform):
     """Maintains an exponential moving average of update.
 
     Args:
@@ -20,12 +20,12 @@ class EMA(Transform):
         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'):
+    def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=True, ema_init: Literal['zeros', 'update'] = 'zeros'):
         defaults = dict(momentum=momentum,dampening=dampening,debiased=debiased,lerp=lerp,ema_init=ema_init)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_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')(settings[0])
@@ -53,8 +53,8 @@ class HeavyBall(EMA):
         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=False, ema_init: Literal['zeros', 'update'] = 'update', target: Target = 'update'):
-        super().__init__(momentum=momentum, dampening=dampening, debiased=debiased, lerp=lerp, ema_init=ema_init, target=target)
+    def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=False, ema_init: Literal['zeros', 'update'] = 'update'):
+        super().__init__(momentum=momentum, dampening=dampening, debiased=debiased, lerp=lerp, ema_init=ema_init)
 
 def nag_(
     tensors_: TensorList,
@@ -74,7 +74,7 @@ def nag_(
     return tensors_
 
 
-class NAG(Transform):
+class NAG(TensorTransform):
     """Nesterov accelerated gradient method (nesterov momentum).
 
     Args:
@@ -84,12 +84,12 @@ class NAG(Transform):
             whether to use linear interpolation, if True, this becomes similar to exponential moving average. Defaults to False.
         target (Target, optional): target to apply EMA to. Defaults to 'update'.
     """
-    def __init__(self, momentum:float=0.9, dampening:float=0, lerp=False, target: Target = 'update'):
+    def __init__(self, momentum:float=0.9, dampening:float=0, lerp=False):
         defaults = dict(momentum=momentum,dampening=dampening, lerp=lerp)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
         lerp = self.settings[params[0]]['lerp']
 

torchzero/modules/ops/__init__.py

@@ -12,8 +12,8 @@ from .binary import (
     CopyMagnitude,
     CopySign,
     Div,
-    Graft,
-    GraftToUpdate,
+    GraftInputToOutput,
+    GraftInputToOutput,
     GramSchimdt,
     Maximum,
     Minimum,
@@ -21,7 +21,7 @@ from .binary import (
     Pow,
     RCopySign,
     RDiv,
-    RGraft,
+    GraftOutputToInput,
     RPow,
     RSub,
     Sub,
@@ -38,7 +38,7 @@ from .higher_level import (
 from .multi import (
     ClipModules,
     DivModules,
-    GraftModules,
+    Graft,
     LerpModules,
     MultiOperationBase,
     PowModules,

torchzero/modules/ops/accumulate.py

@@ -1,90 +1,90 @@
 import torch
 
-from ...core import Target, Transform
+from ...core import  TensorTransform
 from ...utils import TensorList, unpack_states
 
-class AccumulateSum(Transform):
+class AccumulateSum(TensorTransform):
     """Accumulates sum of all past updates.
 
     Args:
         decay (float, optional): decays the accumulator. Defaults to 0.
         target (Target, optional): target. Defaults to 'update'.
     """
-    def __init__(self, decay: float = 0, target: Target = 'update',):
+    def __init__(self, decay: float = 0):
         defaults = dict(decay=decay)
-        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):
         sum = unpack_states(states, tensors, 'sum', cls=TensorList)
         decay = [1-s['decay'] for s in settings]
         return sum.add_(tensors).lazy_mul(decay, clone=True)
 
-class AccumulateMean(Transform):
+class AccumulateMean(TensorTransform):
     """Accumulates mean of all past updates.
 
     Args:
         decay (float, optional): decays the accumulator. Defaults to 0.
         target (Target, optional): target. Defaults to 'update'.
     """
-    def __init__(self, decay: float = 0, target: Target = 'update',):
+    def __init__(self, decay: float = 0):
         defaults = dict(decay=decay)
-        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):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
         mean = unpack_states(states, tensors, 'mean', cls=TensorList)
         decay = [1-s['decay'] for s in settings]
         return mean.add_(tensors).lazy_mul(decay, clone=True).div_(step)
 
-class AccumulateProduct(Transform):
+class AccumulateProduct(TensorTransform):
     """Accumulates product of all past updates.
 
     Args:
         decay (float, optional): decays the accumulator. Defaults to 0.
         target (Target, optional): target. Defaults to 'update'.
     """
-    def __init__(self, decay: float = 0, target: Target = 'update',):
+    def __init__(self, decay: float = 0, target = 'update',):
         defaults = dict(decay=decay)
-        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):
         prod = unpack_states(states, tensors, 'prod', cls=TensorList)
         decay = [1-s['decay'] for s in settings]
         return prod.mul_(tensors).lazy_mul(decay, clone=True)
 
-class AccumulateMaximum(Transform):
+class AccumulateMaximum(TensorTransform):
     """Accumulates maximum of all past updates.
 
     Args:
         decay (float, optional): decays the accumulator. Defaults to 0.
         target (Target, optional): target. Defaults to 'update'.
     """
-    def __init__(self, decay: float = 0, target: Target = 'update',):
+    def __init__(self, decay: float = 0):
         defaults = dict(decay=decay)
-        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):
         maximum = unpack_states(states, tensors, 'maximum', cls=TensorList)
         decay = [1-s['decay'] for s in settings]
         return maximum.maximum_(tensors).lazy_mul(decay, clone=True)
 
-class AccumulateMinimum(Transform):
+class AccumulateMinimum(TensorTransform):
     """Accumulates minimum of all past updates.
 
     Args:
         decay (float, optional): decays the accumulator. Defaults to 0.
         target (Target, optional): target. Defaults to 'update'.
     """
-    def __init__(self, decay: float = 0, target: Target = 'update',):
+    def __init__(self, decay: float = 0):
         defaults = dict(decay=decay)
-        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):
         minimum = unpack_states(states, tensors, 'minimum', cls=TensorList)
         decay = [1-s['decay'] for s in settings]
         return minimum.minimum_(tensors).lazy_mul(decay, clone=True)