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

torchzero/modules/experimental/absoap.py

@@ -1,250 +0,0 @@
-from operator import itemgetter
-from typing import Literal
-
-import torch
-
-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_absoap_covariances_(
-    g1: torch.Tensor,
-    g2: torch.Tensor,
-    GGs_: list[torch.Tensor | None],
-    beta: float | None,
-):
-    for i, GG in enumerate(GGs_):
-        if GG is None: continue
-
-        axes = list(range(i)) + list(range(i + 1, g1.ndim)) # this works fine with 1d params
-        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]
-
-
-Source=Literal['p','g','s','y', 'gy', 'sy', 'sn', 'yn', 'gys', 'sys']
-class ABSOAP(Transform):
-    """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
-
-    """
-    def __init__(
-        self,
-        beta1: float = 0.95,
-        beta2: float = 0.95,
-        shampoo_beta: float | None = 0.95,
-        precond_freq: int = 10,
-        merge_small: bool = True,
-        max_dim: int = 2_000,
-        precondition_1d: bool = True,
-        eps: float = 1e-8,
-        decay: float | None = None,
-        alpha: float = 1,
-        bias_correction: bool = True,
-        scale_by_s: bool = True,
-        gg1: Source='g',
-        gg2: Source='g',
-        ema1: Source='g',
-        ema2: tuple[Source, Source] = ('g','g'),
-        rel1: bool=False,
-        rel2: bool=False,
-        norm: bool = False,
-    ):
-        defaults = dict(
-            beta1=beta1,
-            beta2=beta2,
-            shampoo_beta=shampoo_beta,
-            precond_freq=precond_freq,
-            merge_small=merge_small,
-            max_dim=max_dim,
-            precondition_1d=precondition_1d,
-            eps=eps,
-            decay=decay,
-            bias_correction=bias_correction,
-            alpha=alpha,
-            scale_by_s=scale_by_s,
-            ema1=ema1,
-            ema2=ema2,
-            first=gg1,
-            second=gg2,
-            rel1=rel1, rel2=rel2,
-            norm=norm,
-        )
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        updates = []
-        # update preconditioners
-        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)
-            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)
-
-            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
-
-            p_prev = state['p_prev']
-            g_prev = state['g_prev']
-            s = p - p_prev
-            y = t - g_prev
-
-            # keep malding
-            p_norm = torch.linalg.vector_norm(p) # pylint:disable=not-callable
-            g_norm = torch.linalg.vector_norm(t) # pylint:disable=not-callable
-            s_norm = torch.linalg.vector_norm(s) # pylint:disable=not-callable
-            y_norm = torch.linalg.vector_norm(y) # pylint:disable=not-callable
-
-            sn = p - p_prev * (p_norm / torch.linalg.vector_norm(p_prev))# pylint:disable=not-callable
-            yn = t - g_prev * (g_norm / torch.linalg.vector_norm(g_prev))# pylint:disable=not-callable
-
-            if scale_by_s: y /= s_norm.clip(min=1e-8) # pylint:disable=not-callable
-
-            state['p_prev'].copy_(p)
-            state['g_prev'].copy_(t)
-
-            def _get(c: Source):
-                if c == 'p': return p
-                if c == 'g': return t
-                if c == 's': return s
-                if c == 'y': return y
-                if c == 'sn': return sn
-                if c == 'yn': return yn
-                if c == 'gy': return t+y
-                if c == 'sy': return s+y
-                if c == 'gys':
-                    y_scaled = y * (g_norm/y_norm.clip(min=1e-8))
-                    return t+y_scaled
-                if c == 'sys':
-                    y_scaled = y * (s_norm/y_norm.clip(min=1e-8))
-                    return s+y_scaled
-                raise RuntimeError("Big Chungus")
-
-            t1 = _get(first)
-            if rel1: t1 = t1 * p.abs().clip(min=1e-6)
-            t2 = _get(second)
-            if rel2: t2 = t2 * p.abs().clip(min=1e-6)
-
-            t_ema1 = _get(ema1)
-            t_ema2s = _get(ema2[0]), _get(ema2[1])
-
-            if norm:
-                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.zeros_like(t)
-
-                if not precondition_1d and t.ndim <= 1:
-                    state['GG'] = []
-
-                else:
-                    state['GG'] = [torch.zeros(sh, sh, dtype=t.dtype, device=t.device) if 1<sh<max_dim else None for sh in t.shape]
-
-                # either scalar parameter, 1d with precondition_1d=False, or all dims are too big.
-                if len([i is not None for i in state['GG']]) == 0:
-                    state['GG'] = None
-
-                if state['GG'] is not None:
-                    update_absoap_covariances_(t1, t2, GGs_=state['GG'], beta=shampoo_beta)
-                    state['Q'] = get_orthogonal_matrix(state['GG'])
-
-                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.
-
-            # Projecting gradients to the eigenbases of Shampoo's preconditioner
-            # i.e. projecting to the eigenbases of matrices in state['GG']
-            z1_projected = None
-            z2_projected = None
-
-            if state['GG'] is not None:
-                z1_projected = project(t_ema2s[0], state['Q'])
-                if ema2[0] == ema2[1]: z2_projected = z1_projected
-                else: z2_projected = project(t_ema2s[1], state['Q'])
-
-            # exponential moving averages
-            # this part could be foreached but I will do that at some point its not a big difference compared to preconditioning
-            exp_avg: torch.Tensor = state["exp_avg"]
-            exp_avg_sq: torch.Tensor = state["exp_avg_sq"]
-
-            exp_avg.lerp_(t_ema1, 1-beta1)
-
-            if z1_projected is None:
-                exp_avg_sq.mul_(beta2).addcmul_(*t_ema2s, value=1-beta2)
-            else:
-                assert z2_projected is not None
-                exp_avg_sq.mul_(beta2).addcmul_(z1_projected, z2_projected, value=1-beta2)
-
-            # project exponential moving averages if they are accumulated unprojected
-            exp_avg_projected = exp_avg
-            if z1_projected is not None:
-                exp_avg_projected = project(exp_avg, state['Q'])
-
-            exp_avg_sq_projected = exp_avg_sq
-
-            denom = exp_avg_sq_projected.sqrt().add_(eps)
-            # print(f'{t_projected = }, {exp_avg = }, {exp_avg_projected = }, {exp_avg_sq = }, {exp_avg_sq_projected = }, {denom = }')
-
-            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
-            # to the original space
-            update = exp_avg_projected / denom
-            if z1_projected is not None:
-                update = project_back(update, state["Q"])
-
-            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
-            elif alpha is not None:
-                update *= alpha
-
-            if merge_small:
-                update = _unmerge_small_dims(update, state['flat_sizes'], state['sort_idxs'])
-
-            updates.append(update)
-            state["step"] += 1
-
-            # Update is done after the gradient step to avoid using current gradients in the projection.
-            if state['GG'] is not None:
-                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

@@ -1,112 +0,0 @@
-from operator import itemgetter
-from functools import partial
-
-import torch
-
-from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList
-from ..functional import (
-    debias, debiased_step_size,
-    ema_,
-    sqrt_ema_sq_,
-)
-from ..lr.lr import lazy_lr
-from ..momentum.experimental import sqrt_nag_ema_sq_
-from ..momentum.momentum import nag_
-
-
-def adadam_(
-    tensors: TensorList,
-    exp_avg_: TensorList,
-    exp_avg_sq_: TensorList,
-    exp_avg_qu_: TensorList,
-    alpha: float | NumberList,
-    beta1: float | NumberList,
-    beta2: float | NumberList,
-    precond_beta: float | NumberList,
-    eps: float | NumberList,
-    step: int,
-    pow: float = 2,
-    debiased: bool = True,
-    max_exp_avg_sq_: TensorList | None = None,
-    max_exp_avg_qu_: TensorList | None = None,
-    params_: TensorList | 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)
-    sqrt_exp_avg_qu = sqrt_ema_sq_(tensors/(sqrt_exp_avg_sq+1e-8), exp_avg_sq_=exp_avg_qu_,
-                                   beta=precond_beta,max_exp_avg_sq_=max_exp_avg_qu_, debiased=False,step=step,pow=pow)
-
-    if debiased: alpha = debiased_step_size(step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
-
-    # params is None, return update
-    if params_ is None: return (exp_avg_ / sqrt_exp_avg_qu.add_(eps)).lazy_mul(alpha)
-
-    # update params in-place
-    params_.addcdiv_(exp_avg_, sqrt_exp_avg_qu.add_(eps), -alpha)
-    return None
-
-class Adadam(Module):
-    """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,
-        beta2: float = 0.999,
-        precond_beta: float = 0.999,
-        eps: float = 1e-8,
-        amsgrad: bool = False,
-        alpha: float = 1.,
-        pow: float = 2,
-        debiased: bool = True,
-    ):
-        defaults=dict(beta1=beta1,beta2=beta2,precond_beta=precond_beta,eps=eps,alpha=alpha,amsgrad=amsgrad,pow=pow,debiased=debiased)
-        super().__init__(defaults)
-        self.getter = itemgetter('amsgrad','pow','debiased')
-
-    @torch.no_grad
-    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(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(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(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 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
-
-        var.update = adadam_(
-            tensors=TensorList(var.get_update()),
-            exp_avg_=exp_avg,
-            exp_avg_sq_=exp_avg_sq,
-            exp_avg_qu_=exp_avg_qu,
-            alpha=alpha,
-            beta1=beta1,
-            beta2=beta2,
-            precond_beta=precond_beta,
-            eps=eps,
-            step=step,
-            pow=pow,
-            debiased=debiased,
-            max_exp_avg_sq_=max_exp_avg_sq,
-            max_exp_avg_qu_=max_exp_avg_qu,
-            params_=passed_params,
-        )
-
-        return var

torchzero/modules/experimental/adamY.py

@@ -1,125 +0,0 @@
-from operator import itemgetter
-from functools import partial
-
-import torch
-
-from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList
-from ..functional import (
-    debias, debiased_step_size,
-    ema_,
-    sqrt_ema_sq_,
-)
-from ..lr.lr import lazy_lr
-from ..momentum.experimental import sqrt_nag_ema_sq_
-from ..momentum.momentum import nag_
-
-
-def adamy_(
-    p: TensorList,
-    p_prev: TensorList,
-    g: TensorList,
-    g_prev: TensorList,
-    exp_avg_: TensorList,
-    exp_avg_sq_: TensorList,
-    alpha: float | NumberList,
-    beta1: float | NumberList,
-    beta2: float | NumberList,
-    eps: float | NumberList,
-    step: int,
-    pow: float = 2,
-    debiased: bool = True,
-    max_exp_avg_sq_: TensorList | None = None,
-    params_: TensorList | None = None,
-):
-    """Returns new tensors or updates params in-place."""
-    if step == 1:
-        p_prev.copy_(p)
-        g_prev.copy_(g)
-
-        update = g.clip(-0.1,0.1).lazy_mul_(alpha)
-        if params_ is None: return update
-        params_.sub_(update)
-        return None
-
-    s = p-p_prev
-    y = (g-g_prev).div_(s.global_vector_norm().clip(min=1e-8))
-    p_prev.copy_(p)
-    g_prev.copy_(g)
-
-    exp_avg_ = ema_(g, exp_avg_=exp_avg_, beta=beta1, dampening=0,lerp=True)
-
-    sqrt_exp_avg_sq = sqrt_ema_sq_(y, 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)
-
-    # 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
-
-class AdamY(Module):
-    """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,
-        beta2: float = 0.999,
-        eps: float = 1e-8,
-        amsgrad: bool = False,
-        alpha: float = 1.,
-        pow: float = 2,
-        debiased: bool = True,
-    ):
-        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')
-
-    @torch.no_grad
-    def step(self, var):
-        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
-
-        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(var.params,'exp_avg','exp_avg_sq','max_exp_avg_sq', cls=TensorList)
-        else:
-            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 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(var.params, 'p_prev', cls=TensorList)
-        g_prev = self.get_state(var.params, 'g_prev', cls=TensorList)
-
-
-        var.update = adamy_(
-            p=TensorList(var.params),
-            p_prev=p_prev,
-            g=TensorList(var.get_update()),
-            g_prev=g_prev,
-            exp_avg_=exp_avg,
-            exp_avg_sq_=exp_avg_sq,
-            alpha=alpha,
-            beta1=beta1,
-            beta2=beta2,
-            eps=eps,
-            step=step,
-            pow=pow,
-            debiased=debiased,
-            max_exp_avg_sq_=max_exp_avg_sq,
-            params_=passed_params,
-        )
-
-        return var

torchzero/modules/experimental/adasoap.py

@@ -1,172 +0,0 @@
-from operator import itemgetter
-
-import torch
-
-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_adasoap_covariances_(
-    grad: torch.Tensor,
-    GGs_: list[torch.Tensor | None],
-    GG_sqs: list[torch.Tensor | None],
-    beta: float | None,
-    precond_beta: float | None,
-):
-    for i, (GG, GG_sq) in enumerate(zip(GGs_, GG_sqs)):
-        if GG is None: continue
-        assert GG_sq is not None
-
-        if precond_beta is None: GG_sq.addcmul_(GG, GG)
-        else: GG_sq.mul_(precond_beta).addcmul_(GG, GG, value=1-precond_beta)
-
-        axes = list(range(i)) + list(range(i + 1, grad.ndim)) # this works fine with 1d params
-        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]
-
-
-class AdaSOAP(Transform):
-    """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
-    """
-    def __init__(
-        self,
-        beta1: float = 0.95,
-        beta2: float = 0.95,
-        shampoo_beta: float | None = 0.95,
-        precond_beta: float | None = 0.95,
-        precond_freq: int = 10,
-        merge_small: bool = True,
-        max_dim: int = 2_000,
-        precondition_1d: bool = True,
-        eps: float = 1e-8,
-        decay: float | None = None,
-        alpha: float = 1,
-        unprojected_exp_avg: bool = True,
-        bias_correction: bool = True,
-    ):
-        defaults = dict(
-            beta1=beta1,
-            beta2=beta2,
-            shampoo_beta=shampoo_beta,
-            precond_beta=precond_beta,
-            precond_freq=precond_freq,
-            merge_small=merge_small,
-            max_dim=max_dim,
-            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 apply(self, tensors, params, grads, loss, states, settings):
-        updates = []
-        # update preconditioners
-        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')(setting)
-            precond_beta = setting['precond_beta']
-
-            if merge_small:
-                t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
-
-            # initialize state on 1st step
-            if 'GG' not in state:
-                state["exp_avg"] = torch.zeros_like(t)
-                state["exp_avg_sq"] = torch.zeros_like(t)
-
-                if not precondition_1d and t.ndim <= 1:
-                    state['GG'] = []
-                    state['GG_sq'] = []
-
-                else:
-                    state['GG'] = [torch.zeros(s, s, dtype=t.dtype, device=t.device) if 1<s<max_dim else None for s in t.shape]
-                    state['GG_sq'] = [torch.zeros(s, s, dtype=t.dtype, device=t.device) if 1<s<max_dim else None for s in t.shape]
-
-                # either scalar parameter, 1d with precondition_1d=False, or all dims are too big.
-                if len([i is not None for i in state['GG']]) == 0:
-                    state['GG'] = None
-                    state['GG_sq'] = None
-
-                if state['GG'] is not None:
-                    assert state['GG_sq'] is not None
-                    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 ?
-                # 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']
-            t_projected = None
-            if state['GG'] is not None:
-                t_projected = project(t, state['Q'])
-
-            # exponential moving averages
-            # this part could be foreached but I will do that at some point its not a big difference compared to preconditioning
-            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)
-
-            if t_projected is None:
-                exp_avg_sq.mul_(beta2).addcmul_(t, t, value=1-beta2)
-            else:
-                exp_avg_sq.mul_(beta2).addcmul_(t_projected, t_projected, value=1-beta2)
-
-            # project exponential moving averages if they are accumulated unprojected
-            exp_avg_projected = exp_avg
-            if unprojected_exp_avg and t_projected is not None:
-                exp_avg_projected = project(exp_avg, state['Q'])
-
-            exp_avg_sq_projected = exp_avg_sq
-
-            denom = exp_avg_sq_projected.sqrt().add_(eps)
-            # print(f'{t_projected = }, {exp_avg = }, {exp_avg_projected = }, {exp_avg_sq = }, {exp_avg_sq_projected = }, {denom = }')
-
-            # 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 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
-            elif alpha is not None:
-                update *= alpha
-
-            if merge_small:
-                update = _unmerge_small_dims(update, state['flat_sizes'], state['sort_idxs'])
-
-            updates.append(update)
-            state["step"] += 1
-
-            # Update is done after the gradient step to avoid using current gradients in the projection.
-            if state['GG'] is not None:
-                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'] % setting['precond_freq'] == 0:
-                    state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, GG_precond, state['Q'])
-
-        return updates