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

torchzero/modules/momentum/cautious.py

@@ -48,16 +48,25 @@ class Cautious(Transform):
         eps (float, optional): epsilon for normalization. Defaults to 1e-6.
         mode (str, optional):
             what to do with updates with inconsistent signs.
+            - "zero" - set them to zero (as in paper)
+            - "grad" - set them to the gradient (same as using update magnitude and gradient sign)
+            - "backtrack" - negate them
 
-            "zero" - set them to zero (as in paper)
+    ## Examples:
 
-            "grad" - set them to the gradient
+    Cautious Adam
 
-            "backtrack" - negate them (same as using update magnitude and gradient sign)
+    ```python
+    opt = tz.Modular(
+        bench.parameters(),
+        tz.m.Adam(),
+        tz.m.Cautious(),
+        tz.m.LR(1e-2)
+    )
+    ```
 
-    reference
-        *Cautious Optimizers: Improving Training with One Line of Code.
-        Kaizhao Liang, Lizhang Chen, Bo Liu, Qiang Liu*
+    References:
+        Cautious Optimizers: Improving Training with One Line of Code. Kaizhao Liang, Lizhang Chen, Bo Liu, Qiang Liu
     """
 
     def __init__(
@@ -70,7 +79,7 @@ class Cautious(Transform):
         super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(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)
@@ -89,7 +98,7 @@ class UpdateGradientSignConsistency(Transform):
         super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         normalize, eps = itemgetter('normalize', 'eps')(settings[0])
 
@@ -109,12 +118,9 @@ class IntermoduleCautious(Module):
         eps (float, optional): epsilon for normalization. Defaults to 1e-6.
         mode (str, optional):
             what to do with updates with inconsistent signs.
-
-            "zero" - set them to zero (as in paper)
-
-            "grad" - set them to the gradient
-
-            "backtrack" - negate them (same as using update magnitude and gradient sign)
+            - "zero" - set them to zero (as in paper)
+            - "grad" - set them to the gradient (same as using update magnitude and gradient sign)
+            - "backtrack" - negate them
     """
     def __init__(
         self,
@@ -142,7 +148,7 @@ class IntermoduleCautious(Module):
         compare_var = compare.step(var.clone(clone_update=True))
         var.update_attrs_from_clone_(compare_var)
 
-        mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(self.settings[var.params[0]])
+        mode, normalize, eps = itemgetter('mode', 'normalize', 'eps')(self.defaults)
         var.update = cautious_(
             TensorList(main_var.get_update()),
             TensorList(compare_var.get_update()),
@@ -159,6 +165,18 @@ class ScaleByGradCosineSimilarity(Transform):
 
     Args:
         eps (float, optional): epsilon for division. Defaults to 1e-6.
+
+    ## Examples:
+
+    Scaled Adam
+    ```python
+    opt = tz.Modular(
+        bench.parameters(),
+        tz.m.Adam(),
+        tz.m.ScaleByGradCosineSimilarity(),
+        tz.m.LR(1e-2)
+    )
+    ```
     """
     def __init__(
         self,
@@ -168,12 +186,12 @@ class ScaleByGradCosineSimilarity(Transform):
         super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         eps = settings[0]['eps']
         tensors = TensorList(tensors)
         grads = TensorList(grads)
-        cos_sim = (tensors.dot(grads)) / (tensors.global_vector_norm() * grads.global_vector_norm()).clip(min=eps)
+        cos_sim = tensors.dot(grads) / (tensors.global_vector_norm() * grads.global_vector_norm()).clip(min=eps)
 
         return tensors.mul_(cos_sim)
 
@@ -185,6 +203,20 @@ class ScaleModulesByCosineSimilarity(Module):
         main (Chainable): main module or sequence of modules whose update will be scaled.
         compare (Chainable): module or sequence of modules to compare to
         eps (float, optional): epsilon for division. Defaults to 1e-6.
+
+    ## Examples:
+
+    Adam scaled by similarity to RMSprop
+    ```python
+    opt = tz.Modular(
+        bench.parameters(),
+        tz.m.ScaleModulesByCosineSimilarity(
+            main = tz.m.Adam(),
+            compare = tz.m.RMSprop(0.999, debiased=True),
+        ),
+        tz.m.LR(1e-2)
+    )
+    ```
     """
     def __init__(
         self,
@@ -211,9 +243,9 @@ class ScaleModulesByCosineSimilarity(Module):
 
         m = TensorList(main_var.get_update())
         c = TensorList(compare_var.get_update())
-        eps = self.settings[var.params[0]]['eps']
+        eps = self.defaults['eps']
 
-        cos_sim = (m.dot(c)) / (m.global_vector_norm() * c.global_vector_norm()).clip(min=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

torchzero/modules/momentum/momentum.py

@@ -1,10 +1,44 @@
+from collections import deque
+from operator import itemgetter
 from typing import Literal
 
 import torch
 
 from ...core import Target, Transform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
-from .ema import EMA
+from ..functional import debias, ema_
+
+
+class EMA(Transform):
+    """Maintains an exponential moving average of update.
+
+    Args:
+        momentum (float, optional): momentum (beta). Defaults to 0.9.
+        dampening (float, optional): momentum dampening. Defaults to 0.
+        debiased (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
+        lerp (bool, optional): whether to use linear interpolation. Defaults to True.
+        ema_init (str, optional): initial values for the EMA, "zeros" or "update".
+        target (Target, optional): target to apply EMA to. Defaults to 'update'.
+    """
+    def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=True, ema_init: Literal['zeros', 'update'] = 'zeros', target: Target = 'update'):
+        defaults = dict(momentum=momentum,dampening=dampening,debiased=debiased,lerp=lerp,ema_init=ema_init)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def apply_tensors(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])
+
+        exp_avg = unpack_states(states, tensors, 'exp_avg',
+                                init=torch.zeros_like if ema_init=='zeros' else tensors, cls=TensorList)
+        momentum, dampening = unpack_dicts(settings, 'momentum','dampening', cls=NumberList)
+
+        exp_avg = ema_(TensorList(tensors), exp_avg_=exp_avg,beta=momentum,dampening=dampening,lerp=lerp)
+
+        if debiased: return debias(exp_avg, step=step, beta1=momentum, alpha=1, inplace=False)
+        else: return exp_avg.clone() # this has exp_avg storage so needs to be cloned
+
 
 
 class HeavyBall(EMA):
@@ -55,9 +89,10 @@ class NAG(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
         lerp = self.settings[params[0]]['lerp']
 
         momentum,dampening = unpack_dicts(settings, 'momentum','dampening', cls=NumberList)
         return nag_(TensorList(tensors), velocity_=velocity,momentum=momentum,dampening=dampening,lerp=lerp)
+

torchzero/modules/ops/__init__.py

@@ -7,7 +7,7 @@ from .accumulate import (
 )
 from .binary import (
     Add,
-    BinaryOperation,
+    BinaryOperationBase,
     Clip,
     CopyMagnitude,
     CopySign,
@@ -27,37 +27,20 @@ from .binary import (
     Sub,
     Threshold,
 )
-from .debug import PrintShape, PrintUpdate
-from .misc import (
-    DivByLoss,
-    Dropout,
-    FillLoss,
-    GradientAccumulation,
-    GradSign,
-    GraftGradToUpdate,
-    GraftToGrad,
-    GraftToParams,
-    LastAbsoluteRatio,
-    LastDifference,
-    LastGradDifference,
-    LastProduct,
-    LastRatio,
-    MulByLoss,
-    Multistep,
-    NegateOnLossIncrease,
-    NoiseSign,
-    Previous,
-    Relative,
-    Sequential,
-    UpdateSign,
-    WeightDropout,
+from .higher_level import (
+    CenteredEMASquared,
+    CenteredSqrtEMASquared,
+    Debias,
+    Debias2,
+    EMASquared,
+    SqrtEMASquared,
 )
 from .multi import (
     ClipModules,
     DivModules,
     GraftModules,
     LerpModules,
-    MultiOperation,
+    MultiOperationBase,
     PowModules,
     SubModules,
 )
@@ -66,13 +49,11 @@ from .reduce import (
     Mean,
     MinimumModules,
     Prod,
-    ReduceOperation,
+    ReduceOperationBase,
     Sum,
     WeightedMean,
     WeightedSum,
 )
-from .split import Split
-from .switch import Alternate, Switch
 from .unary import (
     Abs,
     CustomUnaryOperation,
@@ -91,13 +72,12 @@ from .utility import (
     Grad,
     GradToNone,
     Identity,
-    NoOp,
+    Noop,
     Ones,
     Params,
     Randn,
     RandomSample,
     Uniform,
-    Update,
     UpdateToNone,
     Zeros,
 )

torchzero/modules/ops/accumulate.py

@@ -1,11 +1,7 @@
-from collections import deque
-from operator import itemgetter
-from typing import Literal
-
 import torch
 
 from ...core import Target, Transform
-from ...utils import TensorList, NumberList, unpack_states, unpack_dicts
+from ...utils import TensorList, unpack_states
 
 class AccumulateSum(Transform):
     """Accumulates sum of all past updates.
@@ -19,7 +15,7 @@ class AccumulateSum(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(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)
@@ -36,7 +32,7 @@ class AccumulateMean(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(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]
@@ -54,7 +50,7 @@ class AccumulateProduct(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(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)
@@ -71,7 +67,7 @@ class AccumulateMaximum(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(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)
@@ -88,7 +84,7 @@ class AccumulateMinimum(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(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)

torchzero/modules/ops/binary.py

@@ -1,5 +1,4 @@
 #pyright: reportIncompatibleMethodOverride=false
-""""""
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
 from operator import itemgetter
@@ -11,7 +10,7 @@ from ...core import Chainable, Module, Target, Var, maybe_chain
 from ...utils import TensorList, tensorlist
 
 
-class BinaryOperation(Module, ABC):
+class BinaryOperationBase(Module, ABC):
     """Base class for operations that use update as the first operand. This is an abstract class, subclass it and override `transform` method to use it."""
     def __init__(self, defaults: dict[str, Any] | None, **operands: Chainable | Any):
         super().__init__(defaults=defaults)
@@ -47,29 +46,41 @@ class BinaryOperation(Module, ABC):
         return var
 
 
-class Add(BinaryOperation):
+class Add(BinaryOperationBase):
+    """Add :code:`other` to tensors. :code:`other` can be a number or a module.
+
+    If :code:`other` is a module, this calculates :code:`tensors + other(tensors)`
+    """
     def __init__(self, other: Chainable | float, alpha: float = 1):
         defaults = dict(alpha=alpha)
         super().__init__(defaults, other=other)
 
     @torch.no_grad
     def transform(self, var, update: list[torch.Tensor], other: float | list[torch.Tensor]):
-        if isinstance(other, (int,float)): torch._foreach_add_(update, other * self.settings[var.params[0]]['alpha'])
-        else: torch._foreach_add_(update, other, alpha=self.settings[var.params[0]]['alpha'])
+        if isinstance(other, (int,float)): torch._foreach_add_(update, other * self.defaults['alpha'])
+        else: torch._foreach_add_(update, other, alpha=self.defaults['alpha'])
         return update
 
-class Sub(BinaryOperation):
+class Sub(BinaryOperationBase):
+    """Subtract :code:`other` from tensors. :code:`other` can be a number or a module.
+
+    If :code:`other` is a module, this calculates :code:`tensors - other(tensors)`
+    """
     def __init__(self, other: Chainable | float, alpha: float = 1):
         defaults = dict(alpha=alpha)
         super().__init__(defaults, other=other)
 
     @torch.no_grad
     def transform(self, var, update: list[torch.Tensor], other: float | list[torch.Tensor]):
-        if isinstance(other, (int,float)): torch._foreach_sub_(update, other * self.settings[var.params[0]]['alpha'])
-        else: torch._foreach_sub_(update, other, alpha=self.settings[var.params[0]]['alpha'])
+        if isinstance(other, (int,float)): torch._foreach_sub_(update, other * self.defaults['alpha'])
+        else: torch._foreach_sub_(update, other, alpha=self.defaults['alpha'])
         return update
 
-class RSub(BinaryOperation):
+class RSub(BinaryOperationBase):
+    """Subtract tensors from :code:`other`. :code:`other` can be a number or a module.
+
+    If :code:`other` is a module, this calculates :code:`other(tensors) - tensors`
+    """
     def __init__(self, other: Chainable | float):
         super().__init__({}, other=other)
 
@@ -77,7 +88,11 @@ class RSub(BinaryOperation):
     def transform(self, var, update: list[torch.Tensor], other: float | list[torch.Tensor]):
         return other - TensorList(update)
 
-class Mul(BinaryOperation):
+class Mul(BinaryOperationBase):
+    """Multiply tensors by :code:`other`. :code:`other` can be a number or a module.
+
+    If :code:`other` is a module, this calculates :code:`tensors * other(tensors)`
+    """
     def __init__(self, other: Chainable | float):
         super().__init__({}, other=other)
 
@@ -86,7 +101,11 @@ class Mul(BinaryOperation):
         torch._foreach_mul_(update, other)
         return update
 
-class Div(BinaryOperation):
+class Div(BinaryOperationBase):
+    """Divide tensors by :code:`other`. :code:`other` can be a number or a module.
+
+    If :code:`other` is a module, this calculates :code:`tensors / other(tensors)`
+    """
     def __init__(self, other: Chainable | float):
         super().__init__({}, other=other)
 
@@ -95,7 +114,11 @@ class Div(BinaryOperation):
         torch._foreach_div_(update, other)
         return update
 
-class RDiv(BinaryOperation):
+class RDiv(BinaryOperationBase):
+    """Divide :code:`other` by tensors. :code:`other` can be a number or a module.
+
+    If :code:`other` is a module, this calculates :code:`other(tensors) / tensors`
+    """
     def __init__(self, other: Chainable | float):
         super().__init__({}, other=other)
 
@@ -103,7 +126,11 @@ class RDiv(BinaryOperation):
     def transform(self, var, update: list[torch.Tensor], other: float | list[torch.Tensor]):
         return other / TensorList(update)
 
-class Pow(BinaryOperation):
+class Pow(BinaryOperationBase):
+    """Take tensors to the power of :code:`exponent`. :code:`exponent` can be a number or a module.
+
+    If :code:`exponent` is a module, this calculates :code:`tensors ^ exponent(tensors)`
+    """
     def __init__(self, exponent: Chainable | float):
         super().__init__({}, exponent=exponent)
 
@@ -112,7 +139,11 @@ class Pow(BinaryOperation):
         torch._foreach_pow_(update, exponent)
         return update
 
-class RPow(BinaryOperation):
+class RPow(BinaryOperationBase):
+    """Take :code:`other` to the power of tensors. :code:`other` can be a number or a module.
+
+    If :code:`other` is a module, this calculates :code:`other(tensors) ^ tensors`
+    """
     def __init__(self, other: Chainable | float):
         super().__init__({}, other=other)
 
@@ -122,7 +153,11 @@ class RPow(BinaryOperation):
         torch._foreach_pow_(other, update)
         return other
 
-class Lerp(BinaryOperation):
+class Lerp(BinaryOperationBase):
+    """Does a linear interpolation of tensors and :code:`end` module based on a scalar :code:`weight`.
+
+    The output is given by :code:`output = tensors + weight * (end(tensors) - tensors)`
+    """
     def __init__(self, end: Chainable, weight: float):
         defaults = dict(weight=weight)
         super().__init__(defaults, end=end)
@@ -132,7 +167,8 @@ class Lerp(BinaryOperation):
         torch._foreach_lerp_(update, end, weight=self.get_settings(var.params, 'weight'))
         return update
 
-class CopySign(BinaryOperation):
+class CopySign(BinaryOperationBase):
+    """Returns tensors with sign copied from :code:`other(tensors)`."""
     def __init__(self, other: Chainable):
         super().__init__({}, other=other)
 
@@ -140,7 +176,8 @@ class CopySign(BinaryOperation):
     def transform(self, var, update: list[torch.Tensor], other: list[torch.Tensor]):
         return [u.copysign_(o) for u, o in zip(update, other)]
 
-class RCopySign(BinaryOperation):
+class RCopySign(BinaryOperationBase):
+    """Returns :code:`other(tensors)` with sign copied from tensors."""
     def __init__(self, other: Chainable):
         super().__init__({}, other=other)
 
@@ -149,7 +186,11 @@ class RCopySign(BinaryOperation):
         return [o.copysign_(u) for u, o in zip(update, other)]
 CopyMagnitude = RCopySign
 
-class Clip(BinaryOperation):
+class Clip(BinaryOperationBase):
+    """clip tensors to be in  :code:`(min, max)` range. :code:`min` and :code:`max: can be None, numbers or modules.
+
+    If code:`min` and :code:`max`:  are modules, this calculates :code:`tensors.clip(min(tensors), max(tensors))`.
+    """
     def __init__(self, min: float | Chainable | None = None, max: float | Chainable | None = None):
         super().__init__({}, min=min, max=max)
 
@@ -157,8 +198,11 @@ class Clip(BinaryOperation):
     def transform(self, var, update: list[torch.Tensor], min: float | list[torch.Tensor] | None, max: float | list[torch.Tensor] | None):
         return TensorList(update).clamp_(min=min,  max=max)
 
-class MirroredClip(BinaryOperation):
-    """clip by -value, value"""
+class MirroredClip(BinaryOperationBase):
+    """clip tensors to be in  :code:`(-value, value)` range. :code:`value` can be a number or a module.
+
+    If :code:`value` is a module, this calculates :code:`tensors.clip(-value(tensors), value(tensors))`
+    """
     def __init__(self, value: float | Chainable):
         super().__init__({}, value=value)
 
@@ -167,19 +211,19 @@ class MirroredClip(BinaryOperation):
         min = -value if isinstance(value, (int,float)) else [-v for v in value]
         return TensorList(update).clamp_(min=min,  max=value)
 
-class Graft(BinaryOperation):
-    """use direction from update and magnitude from `magnitude` module"""
+class Graft(BinaryOperationBase):
+    """Outputs tensors rescaled to have the same norm as :code:`magnitude(tensors)`."""
     def __init__(self, magnitude: Chainable, tensorwise:bool=True, ord:float=2, eps:float = 1e-6):
         defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps)
         super().__init__(defaults, magnitude=magnitude)
 
     @torch.no_grad
     def transform(self, var, update: list[torch.Tensor], magnitude: list[torch.Tensor]):
-        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(self.settings[var.params[0]])
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(self.defaults)
         return TensorList(update).graft_(magnitude, tensorwise=tensorwise, ord=ord, eps=eps)
 
-class RGraft(BinaryOperation):
-    """use direction from `direction` module and magnitude from update"""
+class RGraft(BinaryOperationBase):
+    """Outputs :code:`magnitude(tensors)` rescaled to have the same norm as tensors"""
 
     def __init__(self, direction: Chainable, tensorwise:bool=True, ord:float=2, eps:float = 1e-6):
         defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps)
@@ -187,12 +231,13 @@ class RGraft(BinaryOperation):
 
     @torch.no_grad
     def transform(self, var, update: list[torch.Tensor], direction: list[torch.Tensor]):
-        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(self.settings[var.params[0]])
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(self.defaults)
         return TensorList(direction).graft_(update, tensorwise=tensorwise, ord=ord, eps=eps)
 
 GraftToUpdate = RGraft
 
-class Maximum(BinaryOperation):
+class Maximum(BinaryOperationBase):
+    """Outputs :code:`maximum(tensors, other(tensors))`"""
     def __init__(self, other: Chainable):
         super().__init__({}, other=other)
 
@@ -201,7 +246,8 @@ class Maximum(BinaryOperation):
         torch._foreach_maximum_(update, other)
         return update
 
-class Minimum(BinaryOperation):
+class Minimum(BinaryOperationBase):
+    """Outputs :code:`minimum(tensors, other(tensors))`"""
     def __init__(self, other: Chainable):
         super().__init__({}, other=other)
 
@@ -211,26 +257,27 @@ class Minimum(BinaryOperation):
         return update
 
 
-class GramSchimdt(BinaryOperation):
-    """makes update orthonormal to `other`"""
+class GramSchimdt(BinaryOperationBase):
+    """outputs tensors made orthogonal to `other(tensors)` via Gram-Schmidt."""
     def __init__(self, other: Chainable):
         super().__init__({}, other=other)
 
     @torch.no_grad
     def transform(self, var, update: list[torch.Tensor], other: list[torch.Tensor]):
         update = TensorList(update); other = TensorList(other)
-        return update - (other*update) / ((other*other) + 1e-8)
+        min = torch.finfo(update[0].dtype).tiny * 2
+        return update - (other*update) / (other*other).clip(min=min)
 
 
-class Threshold(BinaryOperation):
-    """update above/below threshold, value at and below"""
+class Threshold(BinaryOperationBase):
+    """Outputs tensors thresholded such that values above :code:`threshold` are set to :code:`value`."""
     def __init__(self, threshold: Chainable | float, value: Chainable | float, update_above: bool):
         defaults = dict(update_above=update_above)
         super().__init__(defaults, threshold=threshold, value=value)
 
     @torch.no_grad
     def transform(self, var, update: list[torch.Tensor], threshold: list[torch.Tensor] | float, value: list[torch.Tensor] | float):
-        update_above = self.settings[var.params[0]]['update_above']
+        update_above = self.defaults['update_above']
         update = TensorList(update)
         if update_above:
             if isinstance(value, list): return update.where_(update>threshold, value)