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

torchzero/modules/experimental/__init__.py

@@ -3,13 +3,22 @@ from .adadam import Adadam
 from .adamY import AdamY
 from .adasoap import AdaSOAP
 from .curveball import CurveBall
-from .soapy import SOAPY
+from .eigendescent import EigenDescent
+from .etf import (
+    ExponentialTrajectoryFit,
+    ExponentialTrajectoryFitV2,
+    PointwiseExponential,
+)
 from .gradmin import GradMin
+from .newton_solver import NewtonSolver
+from .newtonnewton import NewtonNewton
 from .reduce_outward_lr import ReduceOutwardLR
+from .soapy import SOAPY
 from .spectral import SpectralPreconditioner
+from .structured_newton import StructuredNewton
 from .subspace_preconditioners import (
     HistorySubspacePreconditioning,
     RandomSubspacePreconditioning,
 )
-from .tropical_newton import TropicalNewton
-from .newton_solver import NewtonSolver
+from .tada import TAda
+from .diagonal_higher_order_newton import DiagonalHigherOrderNewton

torchzero/modules/experimental/absoap.py

@@ -1,12 +1,14 @@
 from operator import itemgetter
+from typing import Literal
 
 import torch
-from typing import Literal
-from ...core import Chainable, Transform, apply
+
+from ...core import Chainable, Transform
 from ..optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
+from ..optimizers.soap import project, project_back, get_orthogonal_matrix, get_orthogonal_matrix_QR
 
 @torch.no_grad
-def update_soap_covariances_(
+def update_absoap_covariances_(
     g1: torch.Tensor,
     g2: torch.Tensor,
     GGs_: list[torch.Tensor | None],
@@ -19,138 +21,33 @@ def update_soap_covariances_(
         if beta is None: GG.add_(torch.tensordot(g1, g2, (axes, axes))) # pyright:ignore[reportArgumentType]
         else: GG.lerp_(torch.tensordot(g1, g2, (axes, axes)), 1-beta) # pyright:ignore[reportArgumentType]
 
-@torch.no_grad
-def project(tensors: torch.Tensor, Q: list[torch.Tensor | None]):
-    """
-    Projects the gradient to the eigenbases of the preconditioner.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat, dims=[[0], [0]]) # pyright:ignore[reportArgumentType]
-        else:
-            # I don't understand this part but it is in https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
 
-@torch.no_grad
-def project_back(tensors: torch.Tensor, Q: list[torch.Tensor| None]):
-    """
-    Projects the gradient back to the original space.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat,dims=[[0], [1]]) # pyright:ignore[reportArgumentType]
-        else:
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-@torch.no_grad
-def get_orthogonal_matrix(mat: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using torch.linalg.eigh decomposition.
-    """
-    matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m in mat:
-        if m is None: continue
-        if len(m) == 0:
-            matrix.append([])
-            continue
-        if m.dtype != torch.float:
-            original_type = m.dtype
-            original_device = m.device
-            matrix.append(m.float())
-        else:
-            float_data = True
-            matrix.append(m)
-
-    final = []
-    for m in matrix:
-        if len(m) == 0:
-            final.append([])
-            continue
-        try:
-            _, Q = torch.linalg.eigh(m+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-        except Exception:
-            _, Q = torch.linalg.eigh(m.to(torch.float64)+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-            Q = Q.to(m.dtype)
-        Q = torch.flip(Q, [1])
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-    return final
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py#L240
-@torch.no_grad
-def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | None], Q_list: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using one round of power iteration
-    followed by torch.linalg.qr decomposition.
-    """
-    matrix = []
-    orth_matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m,o in zip(GG, Q_list):
-        if m is None: continue
-        assert o is not None
-
-        if len(m) == 0:
-            matrix.append([])
-            orth_matrix.append([])
-            continue
-        if m.data.dtype != torch.float:
-            original_type = m.data.dtype
-            original_device = m.data.device
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-        else:
-            float_data = True
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-
-    final = []
-    for ind, (m,o) in enumerate(zip(matrix, orth_matrix)):
-        if len(m)==0:
-            final.append([])
-            continue
-        est_eig = torch.diag(o.T @ m @ o)
-        sort_idx = torch.argsort(est_eig, descending=True)
-        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
-        o = o[:,sort_idx]
-        power_iter = m @ o
-        Q, _ = torch.linalg.qr(power_iter) # pylint:disable=not-callable
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-
-    return final, exp_avg_sq
-
-Source=Literal['p','g','s','y', 'gy', 'sy', 'sn', 'yn', 'gys', 'sys','sn', 'yn']
+Source=Literal['p','g','s','y', 'gy', 'sy', 'sn', 'yn', 'gys', 'sys']
 class ABSOAP(Transform):
-    """SOAP but with two extra letters included in its name in order to improve converence
-
-    so what you can do is choose what goes into what ,and that is supposed to be good.
+    """SOAP but with some extra options for testing. Please note that this is experimental and isn't guaranteed to work.
+
+    Args:
+        scale_by_s - whether to scale y by s
+        gg1 - 1st vector into GGᵀ
+        gg2 - 2nd vector into GGᵀ
+        ema1 - vector into 1st momentum
+        ema2 - 2 vectors into 2nd momentum
+        rel1 - if True, multiplies gg1 by params
+        rel2 - same but for gg2
+        norm - if True, gg1 a and gg2 are normalized, and I need to make that into a letter
+
+    letters:
+        p - params
+        g - grad
+        s - param difference
+        y - grad difference
+        gy - g+y
+        sy - s+y
+        sn - s normalized
+        yn - y normalized
+        gys - g + y#g
+        sys - s + y#s
 
-    new args
-
-    scale by s whether to scale gradient differences by parameter differences
-
-    y_to_ema2 whether to use gradient differences for exponential moving average too
-
-    okay I changed these args into another ones
-
-    BASICALLY THIS IS FOR MY EXPERIMENTS
     """
     def __init__(
         self,
@@ -166,8 +63,8 @@ class ABSOAP(Transform):
         alpha: float = 1,
         bias_correction: bool = True,
         scale_by_s: bool = True,
-        first: Source='g',
-        second: Source='g',
+        gg1: Source='g',
+        gg2: Source='g',
         ema1: Source='g',
         ema2: tuple[Source, Source] = ('g','g'),
         rel1: bool=False,
@@ -189,29 +86,27 @@ class ABSOAP(Transform):
             scale_by_s=scale_by_s,
             ema1=ema1,
             ema2=ema2,
-            first=first,
-            second=second,
+            first=gg1,
+            second=gg2,
             rel1=rel1, rel2=rel2,
             norm=norm,
         )
         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]
+        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')(settings)
-            scale_by_s = settings['scale_by_s']
-            ema1 = settings['ema1']
-            ema2 = settings['ema2']
-            first=settings['first']
-            second=settings['second']
-            rel1 = settings['rel1']; rel2 = settings['rel2']
-            norm=settings['norm']
+                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'alpha')(setting)
+            scale_by_s = setting['scale_by_s']
+            ema1 = setting['ema1']
+            ema2 = setting['ema2']
+            first=setting['first']
+            second=setting['second']
+            rel1 = setting['rel1']; rel2 = setting['rel2']
+            norm=setting['norm']
 
             if merge_small:
                 t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
@@ -219,8 +114,8 @@ class ABSOAP(Transform):
             if 'g_prev' not in state:
                 state['p_prev'] = p.clone()
                 state['g_prev'] = t.clone()
-                updates.append(tensors[i].clip(-0.1,0.1))
-                continue
+                # updates.append(tensors[i].clip(-0.1,0.1))
+                # continue
 
             p_prev = state['p_prev']
             g_prev = state['g_prev']
@@ -270,11 +165,10 @@ class ABSOAP(Transform):
                 t1 = t1/torch.linalg.vector_norm(t1).clip(min=1e-8) # pylint:disable=not-callable
                 t2 = t2/torch.linalg.vector_norm(t2).clip(min=1e-8) # pylint:disable=not-callable
 
-
             # initialize state on 1st step
             if 'GG' not in state:
                 state["exp_avg"] = torch.zeros_like(t)
-                state["exp_avg_sq"] = torch.ones_like(t)
+                state["exp_avg_sq"] = torch.zeros_like(t)
 
                 if not precondition_1d and t.ndim <= 1:
                     state['GG'] = []
@@ -287,7 +181,7 @@ class ABSOAP(Transform):
                     state['GG'] = None
 
                 if state['GG'] is not None:
-                    update_soap_covariances_(t1, t2, GGs_=state['GG'], beta=shampoo_beta)
+                    update_absoap_covariances_(t1, t2, GGs_=state['GG'], beta=shampoo_beta)
                     state['Q'] = get_orthogonal_matrix(state['GG'])
 
                 state['step'] = 0
@@ -334,7 +228,7 @@ class ABSOAP(Transform):
             if z1_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
@@ -349,8 +243,8 @@ class ABSOAP(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_(t1, t2, state['GG'], shampoo_beta)
-                if state['step'] % settings['precond_freq'] == 0:
+                update_absoap_covariances_(t1, t2, state['GG'], shampoo_beta)
+                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

torchzero/modules/experimental/adadam.py

@@ -50,7 +50,7 @@ def adadam_(
     return None
 
 class Adadam(Module):
-    """Adam with a diagonally preconditioned preconditioner."""
+    """Adam with a diagonally preconditioned preconditioner. Please note that this is experimental and isn't guaranteed to work."""
     def __init__(
         self,
         beta1: float = 0.9,
@@ -67,31 +67,32 @@ class Adadam(Module):
         self.getter = itemgetter('amsgrad','pow','debiased')
 
     @torch.no_grad
-    def step(self, vars):
+    def step(self, var):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+        params = var.params
 
-        beta1,beta2,precond_beta,eps,alpha=self.get_settings('beta1','beta2','precond_beta','eps','alpha', params=vars.params, cls=NumberList)
-        amsgrad,pow,debiased = self.getter(self.settings[vars.params[0]])
+        beta1,beta2,precond_beta,eps,alpha=self.get_settings(params, 'beta1','beta2','precond_beta','eps','alpha', cls=NumberList)
+        amsgrad,pow,debiased = self.getter(self.settings[var.params[0]])
 
         if amsgrad:
-            exp_avg, exp_avg_sq, exp_avg_qu, max_exp_avg_sq, max_exp_avg_qu = self.get_state('exp_avg','exp_avg_sq', 'exp_avg_qu', 'max_exp_avg_sq', 'max_exp_avg_qu', params=vars.params, cls=TensorList)
+            exp_avg, exp_avg_sq, exp_avg_qu, max_exp_avg_sq, max_exp_avg_qu = self.get_state(params, 'exp_avg','exp_avg_sq', 'exp_avg_qu', 'max_exp_avg_sq', 'max_exp_avg_qu', cls=TensorList)
         else:
-            exp_avg, exp_avg_sq, exp_avg_qu = self.get_state('exp_avg','exp_avg_sq', 'exp_avg_qu', params=vars.params, cls=TensorList)
+            exp_avg, exp_avg_sq, exp_avg_qu = self.get_state(params, 'exp_avg','exp_avg_sq', 'exp_avg_qu', cls=TensorList)
             max_exp_avg_sq = None
             max_exp_avg_qu = 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
+        if var.is_last:
+            if var.last_module_lrs is not None: alpha = alpha * var.last_module_lrs
+            passed_params = TensorList(var.params)
+            var.stop = True
+            var.skip_update = True
 
         else:
             passed_params = None
 
-        vars.update = adadam_(
-            tensors=TensorList(vars.get_update()),
+        var.update = adadam_(
+            tensors=TensorList(var.get_update()),
             exp_avg_=exp_avg,
             exp_avg_sq_=exp_avg_sq,
             exp_avg_qu_=exp_avg_qu,
@@ -108,4 +109,4 @@ class Adadam(Module):
             params_=passed_params,
         )
 
-        return vars
+        return var

torchzero/modules/experimental/adamY.py

@@ -62,17 +62,7 @@ def adamy_(
     return None
 
 class AdamY(Module):
-    """Adam but uses scaled gradient differences for second momentum.
-
-    Args:
-        beta1 (float, optional): momentum. Defaults to 0.9.
-        beta2 (float, optional): second momentum. Defaults to 0.999.
-        eps (float, optional): epsilon. Defaults to 1e-8.
-        alpha (float, optional): learning rate. Defaults to 1.
-        amsgrad (bool, optional): Whether to divide by maximum of EMA of gradient squares instead. Defaults to False.
-        pow (float, optional): power used in second momentum power and root. Defaults to 2.
-        debiased (bool, optional): whether to apply debiasing to momentums based on current step. Defaults to True.
-    """
+    """Adam but uses scaled gradient differences for second momentum. Please note that this is experimental and isn't guaranteed to work."""
     def __init__(
         self,
         beta1: float = 0.9,
@@ -88,36 +78,36 @@ class AdamY(Module):
         self.getter = itemgetter('amsgrad','pow','debiased')
 
     @torch.no_grad
-    def step(self, vars):
+    def step(self, var):
         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=self.get_settings(var.params, 'beta1','beta2','eps','alpha', cls=NumberList)
+        amsgrad,pow,debiased = self.getter(self.settings[var.params[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 = self.get_state(var.params,'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 = self.get_state(var.params, '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
+        if var.is_last:
+            if var.last_module_lrs is not None: alpha = alpha * var.last_module_lrs
+            passed_params = TensorList(var.params)
+            var.stop = True
+            var.skip_update = True
 
         else:
             passed_params = None
 
-        p_prev = self.get_state('p_prev', params=vars.params, cls=TensorList)
-        g_prev = self.get_state('g_prev', params=vars.params, cls=TensorList)
+        p_prev = self.get_state(var.params, 'p_prev', cls=TensorList)
+        g_prev = self.get_state(var.params, 'g_prev', cls=TensorList)
 
 
-        vars.update = adamy_(
-            p=TensorList(vars.params),
+        var.update = adamy_(
+            p=TensorList(var.params),
             p_prev=p_prev,
-            g=TensorList(vars.get_update()),
+            g=TensorList(var.get_update()),
             g_prev=g_prev,
             exp_avg_=exp_avg,
             exp_avg_sq_=exp_avg_sq,
@@ -132,4 +122,4 @@ class AdamY(Module):
             params_=passed_params,
         )
 
-        return vars
+        return var

torchzero/modules/experimental/adasoap.py

@@ -2,11 +2,18 @@ from operator import itemgetter
 
 import torch
 
-from ...core import Chainable, Transform, apply
+from ...core import Chainable, Transform
 from ...modules.optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
+from ..optimizers.soap import (
+    get_orthogonal_matrix,
+    get_orthogonal_matrix_QR,
+    project,
+    project_back,
+)
+
 
 @torch.no_grad
-def update_soap_covariances_(
+def update_adasoap_covariances_(
     grad: torch.Tensor,
     GGs_: list[torch.Tensor | None],
     GG_sqs: list[torch.Tensor | None],
@@ -24,125 +31,9 @@ def update_soap_covariances_(
         if beta is None: GG.add_(torch.tensordot(grad, grad, (axes, axes))) # pyright:ignore[reportArgumentType]
         else: GG.lerp_(torch.tensordot(grad, grad, (axes, axes)), 1-beta) # pyright:ignore[reportArgumentType]
 
-@torch.no_grad
-def project(tensors: torch.Tensor, Q: list[torch.Tensor | None]):
-    """
-    Projects the gradient to the eigenbases of the preconditioner.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat, dims=[[0], [0]]) # pyright:ignore[reportArgumentType]
-        else:
-            # I don't understand this part but it is in https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
-
-@torch.no_grad
-def project_back(tensors: torch.Tensor, Q: list[torch.Tensor| None]):
-    """
-    Projects the gradient back to the original space.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat,dims=[[0], [1]]) # pyright:ignore[reportArgumentType]
-        else:
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-@torch.no_grad
-def get_orthogonal_matrix(mat: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using torch.linalg.eigh decomposition.
-    """
-    matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m in mat:
-        if m is None: continue
-        if len(m) == 0:
-            matrix.append([])
-            continue
-        if m.dtype != torch.float:
-            original_type = m.dtype
-            original_device = m.device
-            matrix.append(m.float())
-        else:
-            float_data = True
-            matrix.append(m)
-
-    final = []
-    for m in matrix:
-        if len(m) == 0:
-            final.append([])
-            continue
-        try:
-            _, Q = torch.linalg.eigh(m+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-        except Exception:
-            _, Q = torch.linalg.eigh(m.to(torch.float64)+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-            Q = Q.to(m.dtype)
-        Q = torch.flip(Q, [1])
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-    return final
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py#L240
-@torch.no_grad
-def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | None], Q_list: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using one round of power iteration
-    followed by torch.linalg.qr decomposition.
-    """
-    matrix = []
-    orth_matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m,o in zip(GG, Q_list):
-        if m is None: continue
-        assert o is not None
-
-        if len(m) == 0:
-            matrix.append([])
-            orth_matrix.append([])
-            continue
-        if m.data.dtype != torch.float:
-            original_type = m.data.dtype
-            original_device = m.data.device
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-        else:
-            float_data = True
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-
-    final = []
-    for ind, (m,o) in enumerate(zip(matrix, orth_matrix)):
-        if len(m)==0:
-            final.append([])
-            continue
-        est_eig = torch.diag(o.T @ m @ o)
-        sort_idx = torch.argsort(est_eig, descending=True)
-        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
-        o = o[:,sort_idx]
-        power_iter = m @ o
-        Q, _ = torch.linalg.qr(power_iter) # pylint:disable=not-callable
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-
-    return final, exp_avg_sq
 
 class AdaSOAP(Transform):
-    """SOAP with diagonally preconditioned GG^Ts
+    """SOAP with diagonally preconditioned GG^Ts. Please note that this is experimental and isn't guaranteed to work.
 
     precond_beta - beta for GG^T squares
     """
@@ -180,15 +71,14 @@ class AdaSOAP(Transform):
         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]
+        for i,(p,t, state, setting) in enumerate(zip(params, tensors, states, settings)):
+
             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)
-            precond_beta = settings['precond_beta']
+                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'unprojected_exp_avg','alpha')(setting)
+            precond_beta = setting['precond_beta']
 
             if merge_small:
                 t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
@@ -213,14 +103,14 @@ class AdaSOAP(Transform):
 
                 if state['GG'] is not None:
                     assert state['GG_sq'] is not None
-                    update_soap_covariances_(t, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
+                    update_adasoap_covariances_(t, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
                     GG_precond = [GG / (GG_sq+1e-8) if GG is not None and GG_sq is not None else None for GG, GG_sq in zip(state['GG'], state['GG_sq'])]
                     state['Q'] = get_orthogonal_matrix(GG_precond)
 
                 state['step'] = 0
                 updates.append(tensors[i].clip(-0.1,0.1))
                 continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
-                # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
+                # that can mess with other modules scaling
 
             # Projecting gradients to the eigenbases of Shampoo's preconditioner
             # i.e. projecting to the eigenbases of matrices in state['GG']
@@ -259,7 +149,7 @@ class AdaSOAP(Transform):
             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
@@ -274,9 +164,9 @@ class AdaSOAP(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, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
+                update_adasoap_covariances_(t, GGs_=state['GG'], GG_sqs=state['GG_sq'], beta=shampoo_beta, precond_beta=precond_beta)
                 GG_precond = [GG / (GG_sq+1e-8) if GG is not None and GG_sq is not None else None for GG, GG_sq in zip(state['GG'], state['GG_sq'])]
-                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, GG_precond, state['Q'])
 
         return updates