torchzero 0.1.8__py3-none-any.whl → 0.3.2__py3-none-any.whl

torchzero/modules/momentum/matrix_momentum.py

@@ -0,0 +1,124 @@
+from typing import Literal
+
+import torch
+
+from ...core import Module, apply
+from ...utils import NumberList, TensorList, as_tensorlist
+from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
+
+class MatrixMomentum(Module):
+    """
+    May be useful for ill conditioned stochastic quadratic objectives but I need to test this.
+    Evaluates hessian vector product on each step (via finite difference or autograd).
+
+    `mu` is supposed to be smaller than (1/largest eigenvalue), otherwise this will be very unstable.
+
+    Orr, Genevieve, and Todd Leen. "Using curvature information for fast stochastic search." Advances in neural information processing systems 9 (1996).
+    """
+    def __init__(self, mu=0.1, beta:float=1, hvp_mode: Literal['autograd', 'forward', 'central'] = 'forward', h=1e-3, hvp_tfm=None):
+        defaults = dict(mu=mu, beta=beta, hvp_mode=hvp_mode, h=h)
+        super().__init__(defaults)
+
+        if hvp_tfm is not None:
+            self.set_child('hvp_tfm', hvp_tfm)
+
+    @torch.no_grad
+    def step(self, vars):
+        assert vars.closure is not None
+        prev_update = self.get_state('prev_update', params=vars.params, cls=TensorList)
+        hvp_mode = self.settings[vars.params[0]]['hvp_mode']
+        h = self.settings[vars.params[0]]['h']
+
+        mu,beta = self.get_settings('mu','beta', params=vars.params, cls=NumberList)
+
+        if hvp_mode == 'autograd':
+            with torch.enable_grad():
+                grad = vars.get_grad(create_graph=True)
+                hvp_ = TensorList(hvp(vars.params, grads=grad, vec=prev_update, allow_unused=True, retain_graph=False)).detach_()
+
+        elif hvp_mode == 'forward':
+            vars.get_grad()
+            l, hvp_ = hvp_fd_forward(vars.closure, vars.params, vec=prev_update, g_0=vars.grad, h=h, normalize=True)
+            if vars.loss_approx is None: vars.loss_approx = l
+
+        elif hvp_mode == 'central':
+            l, hvp_ = hvp_fd_central(vars.closure, vars.params, vec=prev_update, h=h, normalize=True)
+            if vars.loss_approx is None: vars.loss_approx = l
+
+        else:
+            raise ValueError(hvp_mode)
+
+        if 'hvp_tfm' in self.children:
+            hvp_ = TensorList(apply(self.children['hvp_tfm'], hvp_, params=vars.params, grads=vars.grad, vars=vars))
+
+        update = TensorList(vars.get_update())
+
+        hvp_ = as_tensorlist(hvp_)
+        update.add_(prev_update - hvp_*mu)
+        prev_update.set_(update * beta)
+        vars.update = update
+        return vars
+
+
+class AdaptiveMatrixMomentum(Module):
+    """
+    Mu here is estimated as ||s_k||/||y_k||.
+    """
+    def __init__(self, mu_mul:float=1, beta:float=1, eps=1e-4, hvp_mode: Literal['autograd', 'forward', 'central'] = 'forward', h=1e-3, hvp_tfm=None):
+        defaults = dict(mu_mul=mu_mul, beta=beta, hvp_mode=hvp_mode, h=h, eps=eps)
+        super().__init__(defaults)
+
+        if hvp_tfm is not None:
+            self.set_child('hvp_tfm', hvp_tfm)
+
+    @torch.no_grad
+    def step(self, vars):
+        assert vars.closure is not None
+        prev_update, prev_params, prev_grad = self.get_state('prev_update', 'prev_params', 'prev_grad', params=vars.params, cls=TensorList)
+
+        settings = self.settings[vars.params[0]]
+        hvp_mode = settings['hvp_mode']
+        h = settings['h']
+        eps = settings['eps']
+
+        mu_mul, beta = self.get_settings('mu_mul','beta', params=vars.params, cls=NumberList)
+
+        if hvp_mode == 'autograd':
+            with torch.enable_grad():
+                grad = vars.get_grad(create_graph=True)
+                hvp_ = TensorList(hvp(vars.params, grads=grad, vec=prev_update, allow_unused=True, retain_graph=False)).detach_()
+
+        elif hvp_mode == 'forward':
+            vars.get_grad()
+            l, hvp_ = hvp_fd_forward(vars.closure, vars.params, vec=prev_update, g_0=vars.grad, h=h, normalize=True)
+            if vars.loss_approx is None: vars.loss_approx = l
+
+        elif hvp_mode == 'central':
+            l, hvp_ = hvp_fd_central(vars.closure, vars.params, vec=prev_update, h=h, normalize=True)
+            if vars.loss_approx is None: vars.loss_approx = l
+
+        else:
+            raise ValueError(hvp_mode)
+
+        if 'hvp_tfm' in self.children:
+            hvp_ = TensorList(apply(self.children['hvp_tfm'], hvp_, params=vars.params, grads=vars.grad, vars=vars))
+
+        # adaptive part
+        update = TensorList(vars.get_update())
+
+        s_k = vars.params - prev_params
+        prev_params.copy_(vars.params)
+
+        assert vars.grad is not None
+        y_k = vars.grad - prev_grad
+        prev_grad.copy_(vars.grad)
+
+        ada_mu = (s_k.global_vector_norm() / (y_k.global_vector_norm() + eps)) * mu_mul
+
+        # matrix momentum uppdate
+        hvp_ = as_tensorlist(hvp_)
+        update.add_(prev_update - hvp_*ada_mu)
+        prev_update.set_(update * beta)
+        vars.update = update
+        return vars
+

torchzero/modules/momentum/momentum.py

@@ -1,106 +1,43 @@
-from collections import abc
-
-import torch
-
-from ...tensorlist import TensorList
-from ...core import OptimizerModule
-
-def _heavyball_step(ascent, velocity: TensorList, momentum, dampening: TensorList):
-    velocity.mul_(momentum).add_(ascent * (1 - dampening))
-    return velocity.clone()
-
-class HeavyBall(OptimizerModule):
-    """Polyak's (heavy ball) momentum. Exactly matches pytorch SGD `momentum` option.
-
-    Args:
-        decay (float, optional): momentum decay. Defaults to 0.9.
-        dampening (float, optional): momentum dampening. Defaults to 0.
-    """
-    def __init__(self, momentum: float = 0.9, dampening: float = 0, ):
-        defaults = dict(momentum = momentum, dampening = dampening)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        velocity = self.get_state_key('velocity', init = ascent)
-        settings = self.get_all_group_keys()
-        updated_direction = _heavyball_step(ascent, velocity, settings['momentum'], settings['dampening'])
-        return updated_direction
-
-
-def _nesterov_step_(ascent, velocity: TensorList, momentum, dampening,):
-    # update velocity with the ascent direction
-    velocity += ascent
-
-    # decay velocity (this can be moved before previous line for slightly different results)
-    velocity *= momentum
-
-    # update ascent direction with velocity
-    ascent += velocity * (1 - dampening)
-
-
-class NesterovMomentum(OptimizerModule):
-    """Nesterov momentum. Exactly matches pytorch SGD with `nesterov=True`,
-    except this also supports dampening.
-
-    Args:
-        decay (float, optional): momentum decay. Defaults to 0.9.
-        dampening (float, optional): momentum dampening. Defaults to 0.
-    """
-    def __init__(self, decay: float = 0.9, dampening: float = 0, ):
-        defaults = dict(momentum = decay, dampening = dampening)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        velocity = self.get_state_key('velocity')
-        settings = self.get_all_group_keys()
-        _nesterov_step_(ascent, velocity, settings['momentum'], settings['dampening'])
-        return ascent
-
-class GradientAveraging(OptimizerModule):
-    """Averages last 2 gradients (TODO)"""
-    def __init__(self, dampening: float = 0, ):
-        defaults = dict(dampening = dampening)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        velocity = self.get_state_key('velocity')
-        dampening = self.get_group_key('dampening')
-
-        new_direction = ascent + velocity * (1-dampening)
-        velocity.copy_(ascent)
-
-        return new_direction
-
-
-class RandomCoordinateMomentum(OptimizerModule):
-    """Only uses `p` random coordinates of the new update. Other coordinates remain from previous update.
-    This works but I don't know if it is any good.
-
-    Args:
-        p (float, optional): probability to update velocity with a new weigh value. Defaults to 0.1.
-        nesterov (bool, optional): if False, update uses delayed momentum. Defaults to True.
-    """
-    def __init__(self, p: float = 0.1, nesterov=True):
-        defaults = dict(p=p)
-        super().__init__(defaults)
-        self.nesterov = nesterov
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        velocity = self.get_state_key('velocity', init = ascent)
-        settings = self.get_all_group_keys()
-
-        # pick p veclocity indexes to update with the new ascent direction
-        indexes = ascent.bernoulli_like(settings['p']).as_bool()
-
-        if self.nesterov:
-            # update the velocity at those indexes
-            velocity.masked_set_(mask = indexes, value = ascent)
-            return velocity.clone()
-
-        new_ascent = velocity.clone()
-        velocity.masked_set_(mask = indexes, value = ascent)
-        return new_ascent
+from typing import Literal
+
+import torch
+
+from ...core import Target, Transform
+from ...utils import NumberList, TensorList
+from .ema import EMA
+
+
+class HeavyBall(EMA):
+    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 nag_(
+    tensors_: TensorList,
+    velocity_: TensorList,
+    momentum: float | NumberList,
+    dampening: float | NumberList,
+    lerp: bool = False,
+):
+    """Nesterov momentum.
+
+    Returns `tensors_`"""
+    if lerp: velocity_.lerp_(tensors_, 1 - momentum)
+    else: velocity_.add_(tensors_).mul_(momentum)
+
+    tensors_ += velocity_.lazy_mul(1 - dampening)
+
+    return tensors_
+
+
+class NAG(Transform):
+    def __init__(self, momentum:float=0.9, dampening:float=0, lerp=False, target: Target = 'update'):
+        defaults = dict(momentum=momentum,dampening=dampening, lerp=lerp)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        velocity = self.get_state('velocity', params=params, cls=TensorList)
+        lerp = self.settings[params[0]]['lerp']
+
+        momentum,dampening = self.get_settings('momentum','dampening', params=params, cls=NumberList)
+        return nag_(TensorList(tensors), velocity_=velocity,momentum=momentum,dampening=dampening,lerp=lerp)

torchzero/modules/ops/__init__.py

@@ -0,0 +1,103 @@
+from .accumulate import (
+    AccumulateMaximum,
+    AccumulateMean,
+    AccumulateMinimum,
+    AccumulateProduct,
+    AccumulateSum,
+)
+from .binary import (
+    Add,
+    BinaryOperation,
+    Clip,
+    CopyMagnitude,
+    CopySign,
+    Div,
+    Graft,
+    GraftToUpdate,
+    GramSchimdt,
+    Maximum,
+    Minimum,
+    Mul,
+    Pow,
+    RCopySign,
+    RDiv,
+    RGraft,
+    RPow,
+    RSub,
+    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 .multi import (
+    ClipModules,
+    DivModules,
+    GraftModules,
+    LerpModules,
+    MultiOperation,
+    PowModules,
+    SubModules,
+)
+from .reduce import (
+    MaximumModules,
+    Mean,
+    MinimumModules,
+    Prod,
+    ReduceOperation,
+    Sum,
+    WeightedMean,
+    WeightedSum,
+)
+from .split import Split
+from .switch import Alternate, Switch
+from .unary import (
+    Abs,
+    CustomUnaryOperation,
+    Exp,
+    NanToNum,
+    Negate,
+    Reciprocal,
+    Sign,
+    Sqrt,
+    UnaryLambda,
+    UnaryParameterwiseLambda,
+)
+from .utility import (
+    Clone,
+    Fill,
+    Grad,
+    GradToNone,
+    Identity,
+    NoOp,
+    Ones,
+    Params,
+    Randn,
+    RandomSample,
+    Uniform,
+    Update,
+    UpdateToNone,
+    Zeros,
+)

torchzero/modules/ops/accumulate.py

@@ -0,0 +1,65 @@
+from collections import deque
+from operator import itemgetter
+from typing import Literal
+
+import torch
+
+from ...core import Target, Transform
+from ...utils import TensorList, NumberList
+
+class AccumulateSum(Transform):
+    def __init__(self, decay: float = 0, target: Target = 'update',):
+        defaults = dict(decay=decay)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        sum = self.get_state('sum', params=params, cls=TensorList)
+        decay = self.get_settings('decay', params=params, cls=NumberList)
+        return sum.add_(tensors).lazy_mul(1-decay, clone=True)
+
+class AccumulateMean(Transform):
+    def __init__(self, decay: float = 0, target: Target = 'update',):
+        defaults = dict(decay=decay)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+        mean = self.get_state('mean', params=params, cls=TensorList)
+        decay = self.get_settings('decay', params=params, cls=NumberList)
+        return mean.add_(tensors).lazy_mul(1-decay, clone=True).div_(step)
+
+class AccumulateProduct(Transform):
+    def __init__(self, decay: float = 0, target: Target = 'update',):
+        defaults = dict(decay=decay)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        prod = self.get_state('prod', params=params, cls=TensorList)
+        decay = self.get_settings('decay', params=params, cls=NumberList)
+        return prod.mul_(tensors).lazy_mul(1-decay, clone=True)
+
+class AccumulateMaximum(Transform):
+    def __init__(self, decay: float = 0, target: Target = 'update',):
+        defaults = dict(decay=decay)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        maximum = self.get_state('maximum', params=params, cls=TensorList)
+        decay = self.get_settings('decay', params=params, cls=NumberList)
+        return maximum.maximum_(tensors).lazy_mul(1-decay, clone=True)
+
+class AccumulateMinimum(Transform):
+    def __init__(self, decay: float = 0, target: Target = 'update',):
+        defaults = dict(decay=decay)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        minimum = self.get_state('minimum', params=params, cls=TensorList)
+        decay = self.get_settings('decay', params=params, cls=NumberList)
+        return minimum.minimum_(tensors).lazy_mul(1-decay, clone=True)
+

torchzero/modules/ops/binary.py

@@ -0,0 +1,240 @@
+#pyright: reportIncompatibleMethodOverride=false
+""""""
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Sequence
+from operator import itemgetter
+from typing import Any
+
+import torch
+
+from ...core import Chainable, Module, Target, Vars, maybe_chain
+from ...utils import TensorList, tensorlist
+
+
+class BinaryOperation(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)
+
+        self.operands = {}
+        for k,v in operands.items():
+
+            if isinstance(v, (Module, Sequence)):
+                self.set_child(k, v)
+                self.operands[k] = self.children[k]
+            else:
+                self.operands[k] = v
+
+    @abstractmethod
+    def transform(self, vars: Vars, update: list[torch.Tensor], **operands: Any | list[torch.Tensor]) -> Iterable[torch.Tensor]:
+        """applies the operation to operands"""
+        raise NotImplementedError
+
+    @torch.no_grad
+    def step(self, vars: Vars) -> Vars:
+        # 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
+
+        transformed = self.transform(vars, update=vars.get_update(), **processed_operands)
+        vars.update = list(transformed)
+        return vars
+
+
+class Add(BinaryOperation):
+    def __init__(self, other: Chainable | float, alpha: float = 1):
+        defaults = dict(alpha=alpha)
+        super().__init__(defaults, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: float | list[torch.Tensor]):
+        if isinstance(other, (int,float)): torch._foreach_add_(update, other * self.settings[vars.params[0]]['alpha'])
+        else: torch._foreach_add_(update, other, alpha=self.settings[vars.params[0]]['alpha'])
+        return update
+
+class Sub(BinaryOperation):
+    def __init__(self, other: Chainable | float, alpha: float = 1):
+        defaults = dict(alpha=alpha)
+        super().__init__(defaults, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: float | list[torch.Tensor]):
+        if isinstance(other, (int,float)): torch._foreach_sub_(update, other * self.settings[vars.params[0]]['alpha'])
+        else: torch._foreach_sub_(update, other, alpha=self.settings[vars.params[0]]['alpha'])
+        return update
+
+class RSub(BinaryOperation):
+    def __init__(self, other: Chainable | float):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: float | list[torch.Tensor]):
+        return other - TensorList(update)
+
+class Mul(BinaryOperation):
+    def __init__(self, other: Chainable | float):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: float | list[torch.Tensor]):
+        torch._foreach_mul_(update, other)
+        return update
+
+class Div(BinaryOperation):
+    def __init__(self, other: Chainable | float):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: float | list[torch.Tensor]):
+        torch._foreach_div_(update, other)
+        return update
+
+class RDiv(BinaryOperation):
+    def __init__(self, other: Chainable | float):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: float | list[torch.Tensor]):
+        return other / TensorList(update)
+
+class Pow(BinaryOperation):
+    def __init__(self, exponent: Chainable | float):
+        super().__init__({}, exponent=exponent)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], exponent: float | list[torch.Tensor]):
+        torch._foreach_pow_(update, exponent)
+        return update
+
+class RPow(BinaryOperation):
+    def __init__(self, other: Chainable | float):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: float | list[torch.Tensor]):
+        if isinstance(other, (int, float)): return torch._foreach_pow(other, update) # no in-place
+        torch._foreach_pow_(other, update)
+        return other
+
+class Lerp(BinaryOperation):
+    def __init__(self, end: Chainable, weight: float):
+        defaults = dict(weight=weight)
+        super().__init__(defaults, end=end)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], end: list[torch.Tensor]):
+        torch._foreach_lerp_(update, end, weight=self.get_settings('weight',params=vars))
+        return update
+
+class CopySign(BinaryOperation):
+    def __init__(self, other: Chainable):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: list[torch.Tensor]):
+        return [u.copysign_(o) for u, o in zip(update, other)]
+
+class RCopySign(BinaryOperation):
+    def __init__(self, other: Chainable):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: list[torch.Tensor]):
+        return [o.copysign_(u) for u, o in zip(update, other)]
+CopyMagnitude = RCopySign
+
+class Clip(BinaryOperation):
+    def __init__(self, min: float | Chainable | None = None, max: float | Chainable | None = None):
+        super().__init__({}, min=min, max=max)
+
+    @torch.no_grad
+    def transform(self, vars, 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"""
+    def __init__(self, value: float | Chainable):
+        super().__init__({}, value=value)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], value: float | list[torch.Tensor]):
+        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"""
+    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, vars, update: list[torch.Tensor], magnitude: list[torch.Tensor]):
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(self.settings[vars.params[0]])
+        return TensorList(update).graft_(magnitude, tensorwise=tensorwise, ord=ord, eps=eps)
+
+class RGraft(BinaryOperation):
+    """use direction from `direction` module and magnitude from update"""
+
+    def __init__(self, direction: Chainable, tensorwise:bool=True, ord:float=2, eps:float = 1e-6):
+        defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps)
+        super().__init__(defaults, direction=direction)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], direction: list[torch.Tensor]):
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(self.settings[vars.params[0]])
+        return TensorList(direction).graft_(update, tensorwise=tensorwise, ord=ord, eps=eps)
+
+GraftToUpdate = RGraft
+
+class Maximum(BinaryOperation):
+    def __init__(self, other: Chainable):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: list[torch.Tensor]):
+        torch._foreach_maximum_(update, other)
+        return update
+
+class Minimum(BinaryOperation):
+    def __init__(self, other: Chainable):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: list[torch.Tensor]):
+        torch._foreach_minimum_(update, other)
+        return update
+
+
+class GramSchimdt(BinaryOperation):
+    """makes update orthonormal to `other`"""
+    def __init__(self, other: Chainable):
+        super().__init__({}, other=other)
+
+    @torch.no_grad
+    def transform(self, vars, update: list[torch.Tensor], other: list[torch.Tensor]):
+        update = TensorList(update); other = TensorList(other)
+        return update - (other*update) / ((other*other) + 1e-8)
+
+
+class Threshold(BinaryOperation):
+    """update above/below threshold, value at and below"""
+    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, vars, update: list[torch.Tensor], threshold: list[torch.Tensor] | float, value: list[torch.Tensor] | float):
+        update_above = self.settings[vars.params[0]]['update_above']
+        update = TensorList(update)
+        if update_above:
+            if isinstance(value, list): return update.where_(update>threshold, value)
+            return update.masked_fill_(update<=threshold, value)
+
+        if isinstance(value, list): return update.where_(update<threshold, value)
+        return update.masked_fill_(update>=threshold, value)