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

torchzero/modules/optimizers/adam.py

@@ -3,8 +3,8 @@ from functools import partial
 
 import torch
 
-from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList
+from ...core import Module, Target, Transform, apply_transform, Chainable
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 from ..functional import (
     debias, debiased_step_size,
     ema_,
@@ -27,24 +27,25 @@ def adam_(
     pow: float = 2,
     debiased: bool = True,
     max_exp_avg_sq_: TensorList | None = None,
-    params_: TensorList | None = None,
+
+    # inner args
+    inner: Module | None = None,
+    params: list[torch.Tensor] | None = None,
+    grads: list[torch.Tensor] | None = None,
 ):
     """Returns new tensors or updates params in-place."""
-    exp_avg_ = ema_(tensors, exp_avg_=exp_avg_, beta=beta1, dampening=0,lerp=True)
-
     sqrt_exp_avg_sq = sqrt_ema_sq_(tensors, exp_avg_sq_=exp_avg_sq_, beta=beta2, max_exp_avg_sq_=max_exp_avg_sq_,
                                    debiased=False,step=step,pow=pow)
 
-    if debiased: alpha = debiased_step_size(step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
+    if inner is not None:
+        assert params is not None
+        tensors = TensorList(apply_transform(inner, tensors, params=params, grads=grads))
 
-    # params is None, return update
-    if params_ is None: return (exp_avg_ / sqrt_exp_avg_sq.add_(eps)).lazy_mul(alpha)
-
-    # update params in-place
-    params_.addcdiv_(exp_avg_, sqrt_exp_avg_sq.add_(eps), -alpha)
-    return None
+    exp_avg_ = ema_(tensors, exp_avg_=exp_avg_, beta=beta1, dampening=0,lerp=True)
+    if debiased: alpha = debiased_step_size(step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
+    return (exp_avg_ / sqrt_exp_avg_sq.add_(eps)).lazy_mul(alpha)
 
-class Adam(Module):
+class Adam(Transform):
     """Adam. Divides gradient EMA by EMA of gradient squares with debiased step size. This implementation is slightly different from
     pytorch in that debiasing is applied after adding epsilon.
 
@@ -66,36 +67,29 @@ class Adam(Module):
         alpha: float = 1.,
         pow: float = 2,
         debiased: bool = True,
+        inner: Chainable | None = None
     ):
         defaults=dict(beta1=beta1,beta2=beta2,eps=eps,alpha=alpha,amsgrad=amsgrad,pow=pow,debiased=debiased)
-        super().__init__(defaults)
-        self.getter = itemgetter('amsgrad','pow','debiased')
+        super().__init__(defaults, uses_grad=False)
+
+        if inner is not None: self.set_child('inner', inner)
 
     @torch.no_grad
-    def step(self, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        beta1,beta2,eps,alpha=self.get_settings('beta1','beta2','eps','alpha', params=vars.params, cls=NumberList)
-        amsgrad,pow,debiased = self.getter(self.settings[vars.params[0]])
+        beta1,beta2,eps,alpha=unpack_dicts(settings, 'beta1','beta2','eps','alpha', cls=NumberList)
+        amsgrad,pow,debiased = itemgetter('amsgrad','pow','debiased')(settings[0])
 
         if amsgrad:
-            exp_avg, exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg','exp_avg_sq','max_exp_avg_sq', params=vars.params, cls=TensorList)
+            exp_avg, exp_avg_sq, max_exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', cls=TensorList)
         else:
-            exp_avg, exp_avg_sq = self.get_state('exp_avg','exp_avg_sq', params=vars.params, cls=TensorList)
+            exp_avg, exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', cls=TensorList)
             max_exp_avg_sq = None
 
-        # if this is last module, update parameters in-place with slightly more efficient addcdiv_
-        if vars.is_last:
-            if vars.last_module_lrs is not None: alpha = alpha * vars.last_module_lrs
-            passed_params = TensorList(vars.params)
-            vars.stop = True
-            vars.skip_update = True
-
-        else:
-            passed_params = None
 
-        vars.update = adam_(
-            tensors=TensorList(vars.get_update()),
+        return adam_(
+            tensors=TensorList(tensors),
             exp_avg_=exp_avg,
             exp_avg_sq_=exp_avg_sq,
             alpha=alpha,
@@ -106,7 +100,10 @@ class Adam(Module):
             pow=pow,
             debiased=debiased,
             max_exp_avg_sq_=max_exp_avg_sq,
-            params_=passed_params,
-        )
 
-        return vars
+            # inner args
+            inner=self.children.get("inner", None),
+            params=params,
+            grads=grads,
+
+        )

torchzero/modules/optimizers/lion.py

@@ -1,7 +1,7 @@
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 
 
 def lion_(tensors: TensorList, exp_avg_: TensorList, beta1, beta2,):
@@ -28,8 +28,8 @@ class Lion(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        beta1, beta2 = self.get_settings('beta1', 'beta2', params = params, cls=NumberList)
-        exp_avg = self.get_state('ema', params=params, cls=TensorList)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        beta1, beta2 = unpack_dicts(settings, 'beta1', 'beta2', cls=NumberList)
+        exp_avg = unpack_states(states, tensors, 'ema', cls=TensorList)
         return lion_(TensorList(tensors),exp_avg,beta1,beta2)
 

torchzero/modules/optimizers/muon.py

@@ -164,7 +164,7 @@ class Orthogonalize(TensorwiseTransform):
         method (str, optional):
             Newton-Schulz is very fast, SVD is extremely slow but can be slighly more precise.
         target (str, optional):
-            what to set on vars.
+            what to set on var.
     """
     def __init__(self, ns_steps=5, adjust_lr=False, dual_norm_correction=False,
                  method: Literal['newton-schulz', 'svd'] = 'newton-schulz', target:Target='update'):
@@ -172,9 +172,9 @@ class Orthogonalize(TensorwiseTransform):
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def transform(self, tensor, param, grad, vars):
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
         orthogonalize, ns_steps, dual_norm_correction, adjust_lr, method = itemgetter(
-            'orthogonalize', 'ns_steps', 'dual_norm_correction', 'adjust_lr', 'method')(self.settings[param])
+            'orthogonalize', 'ns_steps', 'dual_norm_correction', 'adjust_lr', 'method')(settings)
 
         if not orthogonalize: return tensor
 
@@ -199,7 +199,7 @@ class DualNormCorrection(TensorwiseTransform):
     def __init__(self, target: Target='update'):
         super().__init__({}, uses_grad=True, target=target)
 
-    def transform(self, tensor, param, grad, vars):
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
         assert grad is not None
         if (tensor.ndim >= 2) and (tensor.size(0) > 1) and (tensor.size(1) > 1):
             return _dual_norm_correction(tensor, grad, batch_first=False)
@@ -213,8 +213,8 @@ class MuonAdjustLR(Transform):
         defaults = dict(alpha=alpha)
         super().__init__(defaults=defaults, uses_grad=False, target=target)
 
-    def transform(self, tensors, params, grads, vars):
-        alphas = self.get_settings('alpha', params=params)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        alphas = [s['alpha'] for s in settings]
         tensors_alphas = [(t, adjust_lr_for_muon(a, t.shape)) for t, a in zip(tensors, alphas) if _is_at_least_2d(t)]
         tensors = [i[0] for i in tensors_alphas]
         a = [i[1] for i in alphas]

torchzero/modules/optimizers/orthograd.py

@@ -30,16 +30,15 @@ class OrthoGrad(Transform):
     Args:
         eps (float, optional): epsilon added to the denominator for numerical stability (default: 1e-30)
         renormalize (bool, optional): whether to graft projected gradient to original gradient norm. Defaults to True.
-        target (Target, optional): what to set on vars. Defaults to 'update'.
+        target (Target, optional): what to set on var. Defaults to 'update'.
     """
     def __init__(self, eps: float = 1e-8, renormalize=True, target: Target = 'update'):
         defaults = dict(eps=eps, renormalize=renormalize)
         super().__init__(defaults, uses_grad=False, target=target)
 
-    def transform(self, tensors, params, grads, vars):
-        settings = self.settings[params[0]]
-        eps = settings['eps']
-        renormalize = settings['renormalize']
+    def apply(self, tensors, params, grads, loss, states, settings):
+        eps = settings[0]['eps']
+        renormalize = settings[0]['renormalize']
 
         params = as_tensorlist(params)
         target = as_tensorlist(tensors)

torchzero/modules/optimizers/rmsprop.py

@@ -3,8 +3,8 @@ from typing import Literal
 
 import torch
 
-from ...core import Module, Target, Transform, Chainable, Vars, apply
-from ...utils import NumberList, TensorList
+from ...core import Module, Target, Transform, Chainable, Var, apply_transform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 from ..functional import sqrt_centered_ema_sq_, sqrt_ema_sq_
 
 
@@ -23,7 +23,6 @@ def rmsprop_(
     inner: Module | None = None,
     params: list[torch.Tensor] | None = None,
     grads: list[torch.Tensor] | None = None,
-    vars: Vars | None = None,
 ):
     """returns `tensors_`"""
     if exp_avg_ is not None:
@@ -36,7 +35,7 @@ def rmsprop_(
 
     if inner is not None:
         assert params is not None
-        tensors_ = TensorList(apply(inner, tensors_, params=params, grads=grads, vars=vars))
+        tensors_ = TensorList(apply_transform(inner, tensors_, params=params, grads=grads))
 
     return tensors_.div_(sqrt_exp_avg_sq.add_(eps))
 
@@ -66,21 +65,20 @@ class RMSprop(Transform):
     ):
         defaults = dict(smoothing=smoothing,eps=eps,centered=centered,debiased=debiased,amsgrad=amsgrad,pow=pow,init=init)
         super().__init__(defaults=defaults, uses_grad=False)
-        self.current_step = 0
+
         if inner is not None:
             self.set_child('inner', inner)
 
-    def transform(self, tensors, params, grads, vars):
-        self.current_step += 1
-
-        smoothing,eps = self.get_settings('smoothing', 'eps', params=params, cls=NumberList)
-        centered,debiased,amsgrad,pow,init = itemgetter('centered','debiased','amsgrad','pow','init')(self.settings[params[0]])
+    def apply(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+        smoothing, eps = unpack_dicts(settings, 'smoothing', 'eps', cls=NumberList)
+        centered, debiased, amsgrad, pow, init = itemgetter('centered','debiased','amsgrad','pow','init')(settings[0])
 
-        exp_avg_sq = self.get_state('exp_avg_sq', params=params, cls=TensorList)
-        exp_avg = self.get_state('exp_avg', params=params, cls=TensorList) if centered else None
-        max_exp_avg_sq = self.get_state('max_exp_avg_sq', params=params, cls=TensorList) if amsgrad else None
+        exp_avg_sq = unpack_states(states, tensors, 'exp_avg_sq', cls=TensorList)
+        exp_avg = unpack_states(states, tensors, 'exp_avg', cls=TensorList) if centered else None
+        max_exp_avg_sq = unpack_states(states, tensors, 'max_exp_avg_sq', cls=TensorList) if amsgrad else None
 
-        if init == 'update' and self.current_step == 1:
+        if init == 'update' and step == 1:
             exp_avg_sq.set_([t**2 for t in tensors])
             if exp_avg is not None: exp_avg.set_([t.clone() for t in tensors])
 
@@ -90,7 +88,7 @@ class RMSprop(Transform):
             smoothing=smoothing,
             eps=eps,
             debiased=debiased,
-            step=self.current_step,
+            step=step,
             exp_avg_=exp_avg,
             max_exp_avg_sq_=max_exp_avg_sq,
             pow=pow,
@@ -99,5 +97,4 @@ class RMSprop(Transform):
             inner=self.children.get("inner", None),
             params=params,
             grads=grads,
-            vars=vars,
         )

torchzero/modules/optimizers/rprop.py

@@ -2,7 +2,7 @@
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList, as_tensorlist
+from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states
 
 
 def _bool_ones_like(x):
@@ -161,20 +161,22 @@ class Rprop(Transform):
         alpha: float = 1,
     ):
         defaults = dict(nplus = nplus, nminus = nminus, alpha = alpha, lb = lb, ub = ub, backtrack=backtrack)
-        self.current_step = 0
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        nplus, nminus, lb, ub, alpha = self.get_settings('nplus', 'nminus', 'lb', 'ub', 'alpha', params=params, cls=NumberList)
-        prev, allowed, magnitudes = self.get_state(
+    def apply(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        nplus, nminus, lb, ub, alpha = unpack_dicts(settings, 'nplus', 'nminus', 'lb', 'ub', 'alpha', cls=NumberList)
+        prev, allowed, magnitudes = unpack_states(
+            states, tensors,
             'prev','allowed','magnitudes',
-            params=params,
             init=[torch.zeros_like, _bool_ones_like, torch.zeros_like],
             cls = TensorList,
         )
 
-        target = rprop_(
+        tensors = rprop_(
             tensors_ = as_tensorlist(tensors),
             prev_ = prev,
             allowed_ = allowed,
@@ -184,12 +186,11 @@ class Rprop(Transform):
             lb = lb,
             ub = ub,
             alpha = alpha,
-            backtrack=self.settings[params[0]]['backtrack'],
-            step=self.current_step,
+            backtrack=settings[0]['backtrack'],
+            step=step,
         )
 
-        self.current_step += 1
-        return target
+        return tensors
 
 
 class ScaleLRBySignChange(Transform):
@@ -220,23 +221,25 @@ class ScaleLRBySignChange(Transform):
     ):
         defaults = dict(nplus=nplus, nminus=nminus, alpha=alpha, lb=lb, ub=ub, use_grad=use_grad)
         super().__init__(defaults, uses_grad=use_grad, target=target)
-        self.current_step = 0
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        target = as_tensorlist(tensors)
-        use_grad = self.settings[params[0]]['use_grad']
+    def apply(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        tensors = as_tensorlist(tensors)
+        use_grad = settings[0]['use_grad']
         if use_grad: cur = as_tensorlist(grads)
-        else: cur = target
+        else: cur = tensors
 
-        nplus, nminus, lb, ub = self.get_settings('nplus', 'nminus', 'lb', 'ub', params=params, cls=NumberList)
-        prev, lrs = self.get_state('prev', 'lrs', params=params, cls=TensorList)
+        nplus, nminus, lb, ub = unpack_dicts(settings, 'nplus', 'nminus', 'lb', 'ub', cls=NumberList)
+        prev, lrs = unpack_states(states, tensors, 'prev', 'lrs', cls=TensorList)
 
-        if self.current_step == 0:
-            lrs.set_(target.full_like(self.get_settings('alpha', params=params)))
+        if step == 0:
+            lrs.set_(tensors.full_like([s['alpha'] for s in settings]))
 
-        target = scale_by_sign_change_(
-            tensors_ = target,
+        tensors = scale_by_sign_change_(
+            tensors_ = tensors,
             cur = cur,
             prev_ = prev,
             lrs_ = lrs,
@@ -244,10 +247,9 @@ class ScaleLRBySignChange(Transform):
             nminus = nminus,
             lb = lb,
             ub = ub,
-            step = self.current_step,
+            step = step,
         )
-        self.current_step += 1
-        return target
+        return tensors
 
 class BacktrackOnSignChange(Transform):
     """Negates or undoes update for parameters where where gradient or update sign changes.
@@ -268,28 +270,28 @@ class BacktrackOnSignChange(Transform):
     def __init__(self, use_grad = False, backtrack = True, target: Target = 'update'):
         defaults = dict(use_grad=use_grad, backtrack=backtrack, target=target)
         super().__init__(defaults, uses_grad=use_grad)
-        self.current_step = 0
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        target = as_tensorlist(tensors)
-        settings = self.settings[params[0]]
-        use_grad = settings['use_grad']
-        backtrack = settings['backtrack']
+    def apply(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        tensors = as_tensorlist(tensors)
+        use_grad = settings[0]['use_grad']
+        backtrack = settings[0]['backtrack']
 
         if use_grad: cur = as_tensorlist(grads)
-        else: cur = target
+        else: cur = tensors
 
-        target = backtrack_on_sign_change_(
-            tensors_ = target,
+        tensors = backtrack_on_sign_change_(
+            tensors_ = tensors,
             cur = cur,
-            prev_ = self.get_state('prev', params=params, cls=TensorList),
+            prev_ = unpack_states(states, tensors, 'prev', cls=TensorList),
             backtrack = backtrack,
-            step = self.current_step,
+            step = step,
         )
 
-        self.current_step += 1
-        return target
+        return tensors
 
 class SignConsistencyMask(Transform):
     """0 if sign changed 1 otherwise"""
@@ -297,10 +299,10 @@ class SignConsistencyMask(Transform):
         super().__init__({}, uses_grad=False, target = target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        prev = self.get_state('prev', params=params, cls=TensorList)
+    def apply(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', cls=TensorList)
         mask = prev.mul_(tensors).gt_(0)
-        prev.set_(tensors)
+        prev.copy_(tensors)
         return mask
 
 
@@ -317,16 +319,18 @@ class SignConsistencyLRs(Transform):
     ):
         defaults = dict(nplus = nplus, nminus = nminus, alpha = alpha, lb = lb, ub = ub)
         super().__init__(defaults, uses_grad=False, target = target)
-        self.current_step = 0
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
         target = as_tensorlist(tensors)
-        nplus, nminus, lb, ub = self.get_settings('nplus', 'nminus', 'lb', 'ub', params=params, cls=NumberList)
-        prev, lrs = self.get_state('prev', 'lrs', params=params, cls=TensorList)
+        nplus, nminus, lb, ub = unpack_dicts(settings, 'nplus', 'nminus', 'lb', 'ub', cls=NumberList)
+        prev, lrs = unpack_states(states, tensors, 'prev', 'lrs', cls=TensorList)
 
-        if self.current_step == 0:
-            lrs.set_(target.full_like(self.get_settings('alpha', params=params)))
+        if step == 0:
+            lrs.set_(target.full_like([s['alpha'] for s in settings]))
 
         target = sign_consistency_lrs_(
             tensors = target,
@@ -336,7 +340,6 @@ class SignConsistencyLRs(Transform):
             nminus = nminus,
             lb = lb,
             ub = ub,
-            step = self.current_step,
+            step = step,
         )
-        self.current_step += 1
         return target.clone()

torchzero/modules/optimizers/shampoo.py

@@ -4,7 +4,7 @@ from functools import partial
 import numpy as np
 import torch
 
-from ...core import Chainable, Transform, apply
+from ...core import Chainable, Transform, apply_transform
 from ...utils.linalg import matrix_power_eigh
 from ...utils import set_storage_
 
@@ -106,7 +106,6 @@ class Shampoo(Transform):
         self,
         decay: float | None = None,
         beta: float | None = None,
-        reg: float = 1e-6,
         update_freq: int = 10,
         exp_override: int | None = None,
         merge_small: bool = True,
@@ -115,25 +114,24 @@ class Shampoo(Transform):
         adagrad_eps: float = 1e-8,
         inner: Chainable | None = None,
     ):
-        defaults = dict(decay=decay, beta=beta, reg=reg, update_freq=update_freq, exp_override=exp_override, merge_small=merge_small, max_dim=max_dim, precondition_1d=precondition_1d,adagrad_eps=adagrad_eps)
+        defaults = dict(decay=decay, beta=beta, update_freq=update_freq, exp_override=exp_override, merge_small=merge_small, max_dim=max_dim, precondition_1d=precondition_1d,adagrad_eps=adagrad_eps)
         super().__init__(defaults, uses_grad=False)
 
         if inner is not None:
             self.set_child('inner', inner)
 
-    def transform(self, tensors, params, grads, vars):
-        merged_target = [] # target with merged dims
+    def apply(self, tensors, params, grads, loss, states, settings):
+        merged_tensors = [] # target with merged dims
 
         # update preconditioners
-        for i,(p,t) in enumerate(zip(params, tensors)):
-            state = self.state[p]
-            settings = self.settings[p]
-            beta, reg, update_freq, exp_override, merge_small, max_dim, precondition_1d = itemgetter(
-                'beta', 'reg', 'update_freq', 'exp_override', 'merge_small', 'max_dim', 'precondition_1d')(settings)
+        for i,(t,state, setting) in enumerate(zip(tensors, states, settings)):
+            beta, update_freq, exp_override, merge_small, max_dim, precondition_1d = itemgetter(
+                'beta', 'update_freq', 'exp_override', 'merge_small', 'max_dim', 'precondition_1d')(setting)
 
             if merge_small:
                 t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
-            merged_target.append(t)
+
+            merged_tensors.append(t)
 
             # initialize accumulators and preconditioners for each dim on 1st step
             if 'accumulators' not in state:
@@ -167,22 +165,18 @@ class Shampoo(Transform):
 
         # inner step
         if 'inner' in self.children:
-            tensors = apply(self.children['inner'], tensors, params=params, grads=grads, vars=vars)
+            tensors = apply_transform(self.children['inner'], tensors, params=params, grads=grads)
 
             # have to merge small dims again
-            merged_target = [] # target with merged dims
-            for i,(p,t) in enumerate(zip(params, tensors)):
-                state = self.state[p]
-                settings = self.settings[p]
-                if settings['merge_small']:
-                    t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, settings['max_dim'])
-                merged_target.append(t)
+            merged_tensors = [] # target with merged dims
+            for i,(t,state, setting) in enumerate(zip(tensors, states, settings)):
+                if setting['merge_small']:
+                    t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, setting['max_dim'])
+                merged_tensors.append(t)
 
         # precondition
-        for i, (p, t) in enumerate(zip(params, merged_target)):
-            state = self.state[p]
-            settings = self.settings[p]
-            decay, merge_small, adagrad_eps= itemgetter('decay', 'merge_small', 'adagrad_eps')(settings)
+        for i,(t,state, setting) in enumerate(zip(merged_tensors, states, settings)):
+            decay, merge_small, adagrad_eps= itemgetter('decay', 'merge_small', 'adagrad_eps')(setting)
 
             if 'diagonal_accumulator' in state:
                 tensors[i] = apply_diagonal_(t, state['diagonal_accumulator'], decay=decay, eps=adagrad_eps)

torchzero/modules/optimizers/soap.py

@@ -2,7 +2,7 @@ from operator import itemgetter
 
 import torch
 
-from ...core import Chainable, Transform, apply
+from ...core import Chainable, Transform, apply_transform
 from ...modules.optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
 
 @torch.no_grad
@@ -152,9 +152,8 @@ class SOAP(Transform):
             epsilon for dividing first momentum by second. Defaults to 1e-8.
         decay (float | None, optional):
             Decays covariance matrix accumulators, this may be useful if `shampoo_beta` is None. Defaults to None.
-        unprojected_exp_avg (bool, optional):
-            whether to update first momentum in unprojected space. Both true and false work and lead to different
-            results but True usually works better. Defaults to True.
+        alpha (float, optional):
+            learning rate. Defaults to 1.
         bias_correction (bool, optional):
             enables adam bias correction. Defaults to True.
     """
@@ -170,7 +169,6 @@ class SOAP(Transform):
         eps: float = 1e-8,
         decay: float | None = None,
         alpha: float = 1,
-        unprojected_exp_avg: bool = True,
         bias_correction: bool = True,
     ):
         defaults = dict(
@@ -183,21 +181,18 @@ class SOAP(Transform):
             precondition_1d=precondition_1d,
             eps=eps,
             decay=decay,
-            unprojected_exp_avg=unprojected_exp_avg,
             bias_correction=bias_correction,
             alpha=alpha,
         )
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         updates = []
         # update preconditioners
-        for i,(p,t) in enumerate(zip(params, tensors)):
-            state = self.state[p]
-            settings = self.settings[p]
-            beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps, unprojected_exp_avg,alpha = itemgetter(
-                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'unprojected_exp_avg','alpha')(settings)
+        for i,(p,t, state, setting) in enumerate(zip(params, tensors, states, settings)):
+            beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps,alpha = itemgetter(
+                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps','alpha')(setting)
 
             if merge_small:
                 t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
@@ -237,10 +232,7 @@ class SOAP(Transform):
             exp_avg: torch.Tensor = state["exp_avg"]
             exp_avg_sq: torch.Tensor = state["exp_avg_sq"]
 
-            if unprojected_exp_avg or t_projected is None:
-                exp_avg.lerp_(t, 1-beta1)
-            else:
-                exp_avg.lerp_(t_projected, 1-beta1)
+            exp_avg.lerp_(t, 1-beta1)
 
             if t_projected is None:
                 exp_avg_sq.mul_(beta2).addcmul_(t, t, value=1-beta2)
@@ -249,7 +241,7 @@ class SOAP(Transform):
 
             # project exponential moving averages if they are accumulated unprojected
             exp_avg_projected = exp_avg
-            if unprojected_exp_avg and t_projected is not None:
+            if t_projected is not None:
                 exp_avg_projected = project(exp_avg, state['Q'])
 
             exp_avg_sq_projected = exp_avg_sq
@@ -260,10 +252,11 @@ class SOAP(Transform):
             # Projecting back the preconditioned (by Adam) exponential moving average of gradients
             # to the original space
             update = exp_avg_projected / denom
+
             if t_projected is not None:
                 update = project_back(update, state["Q"])
 
-            if settings['bias_correction']:
+            if setting['bias_correction']:
                 bias_correction1 = 1.0 - beta1 ** (state["step"]+1)
                 bias_correction2 = 1.0 - beta2 ** (state["step"]+1)
                 update *= ((bias_correction2 ** .5) / bias_correction1) * alpha
@@ -279,7 +272,7 @@ class SOAP(Transform):
             # Update is done after the gradient step to avoid using current gradients in the projection.
             if state['GG'] is not None:
                 update_soap_covariances_(t, state['GG'], shampoo_beta)
-                if state['step'] % settings['precond_freq'] == 0:
+                if state['step'] % setting['precond_freq'] == 0:
                     state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, state['GG'], state['Q'])
 
         return updates