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

torchzero/modules/experimental/diagonal_higher_order_newton.py

@@ -1,225 +0,0 @@
-import itertools
-import math
-import warnings
-from collections.abc import Callable
-from contextlib import nullcontext
-from functools import partial
-from typing import Any, Literal
-
-import numpy as np
-import scipy.optimize
-import torch
-
-from ...core import Chainable, Module, apply_transform
-from ...utils import TensorList, vec_to_tensors, vec_to_tensors_
-from ...utils.derivatives import (
-    hessian_list_to_mat,
-    jacobian_wrt,
-    hvp,
-)
-
-def _poly_eval_diag(s: np.ndarray, c, derivatives):
-    val = float(c) + (derivatives[0] * s).sum(-1)
-
-    if len(derivatives) > 1:
-        for i, d_diag in enumerate(derivatives[1:], 2):
-            val += (d_diag * (s**i)).sum(-1) / math.factorial(i)
-
-    return val
-
-def _proximal_poly_v_diag(x: np.ndarray, c, prox, x0: np.ndarray, derivatives):
-    """Computes the value of the proximal polynomial approximation."""
-    if x.ndim == 2: x = x.T
-    s = x - x0
-
-    val = _poly_eval_diag(s, c, derivatives)
-
-    penalty = 0
-    if prox != 0:
-        penalty = (prox / 2) * (s**2).sum(-1)
-
-    return val + penalty
-
-def _proximal_poly_g_diag(x: np.ndarray, c, prox, x0: np.ndarray, derivatives):
-    """Computes the gradient of the proximal polynomial approximation."""
-    s = x - x0
-
-    g = derivatives[0].copy()
-
-    if len(derivatives) > 1:
-        for i, d_diag in enumerate(derivatives[1:], 2):
-            g += d_diag * (s**(i - 1)) / math.factorial(i - 1)
-
-    if prox != 0:
-        g += prox * s
-
-    return g
-
-def _proximal_poly_H_diag(x: np.ndarray, c, prox, x0: np.ndarray, derivatives):
-    """Computes the Hessian of the proximal polynomial approximation."""
-    s = x - x0
-    n = x.shape[0]
-
-    if len(derivatives) < 2:
-        H_diag = np.zeros(n, dtype=s.dtype)
-    else:
-        H_diag = derivatives[1].copy()
-
-    if len(derivatives) > 2:
-        for i, d_diag in enumerate(derivatives[2:], 3):
-            H_diag += d_diag * (s**(i - 2)) / math.factorial(i - 2)
-
-    if prox != 0:
-        H_diag += prox
-
-    return np.diag(H_diag)
-
-def _poly_minimize(trust_region, prox, de_iters: Any, c, x: torch.Tensor, derivatives):
-    derivatives = [T.detach().cpu().numpy().astype(np.float64) for T in derivatives]
-    x0 = x.detach().cpu().numpy().astype(np.float64) # taylor series center
-    bounds = None
-    if trust_region is not None: bounds = list(zip(x0 - trust_region, x0 + trust_region))
-
-    # if len(derivatives) is 1, only gradient is available, I use that to test proximal penalty and bounds
-    if bounds is None:
-        if len(derivatives) == 1: method = 'bfgs'
-        else: method = 'trust-exact'
-    else:
-        if len(derivatives) == 1: method = 'l-bfgs-b'
-        else: method = 'trust-constr'
-
-    x_init = x0.copy()
-    v0 = _proximal_poly_v_diag(x0, c, prox, x0, derivatives)
-    if de_iters is not None and de_iters != 0:
-        if de_iters == -1: de_iters = None # let scipy decide
-        res = scipy.optimize.differential_evolution(
-            _proximal_poly_v_diag,
-            bounds if bounds is not None else list(zip(x0 - 10, x0 + 10)),
-            args=(c, prox, x0.copy(), derivatives),
-            maxiter=de_iters,
-            vectorized=True,
-        )
-        if res.fun < v0: x_init = res.x
-
-    res = scipy.optimize.minimize(
-        _proximal_poly_v_diag,
-        x_init,
-        method=method,
-        args=(c, prox, x0.copy(), derivatives),
-        jac=_proximal_poly_g_diag,
-        hess=_proximal_poly_H_diag,
-        bounds=bounds
-    )
-
-    return torch.from_numpy(res.x).to(x), res.fun
-
-
-
-class DiagonalHigherOrderNewton(Module):
-    """
-    Hvp with ones doesn't give you the diagonal unless derivatives are diagonal, but somehow it still works,
-    except it doesn't work in all cases except ones where it works.
-    """
-    def __init__(
-        self,
-        order: int = 4,
-        trust_method: Literal['bounds', 'proximal', 'none'] | None = 'bounds',
-        increase: float = 1.5,
-        decrease: float = 0.75,
-        trust_init: float | None = None,
-        trust_tol: float = 1,
-        de_iters: int | None = None,
-        vectorize: bool = True,
-    ):
-        if trust_init is None:
-            if trust_method == 'bounds': trust_init = 1
-            else: trust_init = 0.1
-
-        defaults = dict(order=order, trust_method=trust_method, increase=increase, decrease=decrease, trust_tol=trust_tol, trust_init=trust_init, vectorize=vectorize, de_iters=de_iters)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
-
-        settings = self.settings[params[0]]
-        order = settings['order']
-        increase = settings['increase']
-        decrease = settings['decrease']
-        trust_tol = settings['trust_tol']
-        trust_init = settings['trust_init']
-        trust_method = settings['trust_method']
-        de_iters = settings['de_iters']
-
-        trust_value = self.global_state.get('trust_value', trust_init)
-
-
-        # ------------------------ calculate grad and hessian ------------------------ #
-        with torch.enable_grad():
-            loss = var.loss = var.loss_approx = closure(False)
-
-            g = torch.autograd.grad(loss, params, create_graph=True)
-            var.grad = list(g)
-
-            derivatives = [g]
-            T = g # current derivatives tensor diagonal
-            ones = [torch.ones_like(t) for t in g]
-
-            # get all derivatives up to order
-            for o in range(2, order + 1):
-                T = hvp(params, T, ones, create_graph=o != order)
-                derivatives.append(T)
-
-        x0 = torch.cat([p.ravel() for p in params])
-
-        if trust_method is None: trust_method = 'none'
-        else: trust_method = trust_method.lower()
-
-        if trust_method == 'none':
-            trust_region = None
-            prox = 0
-
-        elif trust_method == 'bounds':
-            trust_region = trust_value
-            prox = 0
-
-        elif trust_method == 'proximal':
-            trust_region = None
-            prox = 1 / trust_value
-
-        else:
-            raise ValueError(trust_method)
-
-        x_star, expected_loss = _poly_minimize(
-            trust_region=trust_region,
-            prox=prox,
-            de_iters=de_iters,
-            c=loss.item(),
-            x=x0,
-            derivatives=[torch.cat([t.ravel() for t in d]) for d in derivatives],
-        )
-
-        # trust region
-        if trust_method != 'none':
-            expected_reduction = loss - expected_loss
-
-            vec_to_tensors_(x_star, params)
-            loss_star = closure(False)
-            vec_to_tensors_(x0, params)
-            reduction = loss - loss_star
-
-            # failed step
-            if reduction <= 0:
-                x_star = x0
-                self.global_state['trust_value'] = trust_value * decrease
-
-            # very good step
-            elif expected_reduction / reduction <= trust_tol:
-                self.global_state['trust_value'] = trust_value * increase
-
-        difference = vec_to_tensors(x0 - x_star, params)
-        var.update = list(difference)
-        return var
-

torchzero/modules/experimental/eigendescent.py

@@ -1,117 +0,0 @@
-from contextlib import nullcontext
-import warnings
-from collections.abc import Callable
-from functools import partial
-import itertools
-from typing import Literal
-
-import torch
-
-from ...core import Chainable, Module, apply_transform
-from ...utils import TensorList, vec_to_tensors
-from ...utils.derivatives import (
-    hessian_list_to_mat,
-    jacobian_wrt, jacobian_and_hessian_wrt, hessian_mat,
-)
-
-def _batched_dot(x, y):
-    return (x.unsqueeze(-2) @ y.unsqueeze(-1)).squeeze(-1).squeeze(-1)
-
-def _cosine_similarity(x, y):
-    denom = torch.linalg.vector_norm(x, dim=-1) * torch.linalg.vector_norm(y, dim=-1).clip(min=torch.finfo(x.dtype).eps) # pylint:disable=not-callable
-    return _batched_dot(x, y) / denom
-
-class EigenDescent(Module):
-    """
-    Uses eigenvectors corresponding to certain eigenvalues. Please note that this is experimental and isn't guaranteed to work.
-
-    Args:
-        mode (str, optional):
-            - largest - use largest eigenvalue unless all eigenvalues are negative, then smallest is used.
-            - smallest - use smallest eigenvalue unless all eigenvalues are positive, then largest is used.
-            - mean-sign - use mean of eigenvectors multiplied by 1 or -1 if they point in opposite direction from gradient.
-            - mean-dot - use mean of eigenvectors multiplied by dot product with gradient.
-            - mean-cosine - use mean of eigenvectors multiplied by cosine similarity with gradient.
-            - mm - for testing.
-
-            Defaults to 'mean-sign'.
-        hessian_method (str, optional): how to calculate hessian. Defaults to "autograd".
-        vectorize (bool, optional): how to calculate hessian. Defaults to True.
-
-    """
-    def __init__(
-        self,
-        mode: Literal['largest', 'smallest','magnitude', 'mean-sign', 'mean-dot', 'mean-cosine', 'mm'] = 'mean-sign',
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
-    ):
-        defaults = dict(hessian_method=hessian_method, vectorize=vectorize, mode=mode)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
-
-        settings = self.settings[params[0]]
-        mode = settings['mode']
-        hessian_method = settings['hessian_method']
-        vectorize = settings['vectorize']
-
-        # ------------------------ calculate grad and hessian ------------------------ #
-        if hessian_method == 'autograd':
-            with torch.enable_grad():
-                loss = var.loss = var.loss_approx = closure(False)
-                g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
-                g_list = [t[0] for t in g_list] # remove leading dim from loss
-                var.grad = g_list
-                H = hessian_list_to_mat(H_list)
-
-        elif hessian_method in ('func', 'autograd.functional'):
-            strat = 'forward-mode' if vectorize else 'reverse-mode'
-            with torch.enable_grad():
-                g_list = var.get_grad(retain_graph=True)
-                H: torch.Tensor = hessian_mat(partial(closure, backward=False), params,
-                                method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
-
-        else:
-            raise ValueError(hessian_method)
-
-
-        # ----------------------------------- solve ---------------------------------- #
-        g = torch.cat([t.ravel() for t in g_list])
-        L, Q = torch.linalg.eigh(H) # L is sorted # pylint:disable=not-callable
-        if mode == 'largest':
-            # smallest eigenvalue if all eigenvalues are negative else largest
-            if L[-1] <= 0: d = Q[0]
-            else: d = Q[-1]
-
-        elif mode == 'smallest':
-            # smallest eigenvalue if negative eigenvalues exist else largest
-            if L[0] <= 0: d = Q[0]
-            else: d = Q[-1]
-
-        elif mode == 'magnitude':
-            # largest by magnitude
-            if L[0].abs() > L[-1].abs(): d = Q[0]
-            else: d = Q[-1]
-
-        elif mode == 'mean-dot':
-            d = ((g.unsqueeze(0) @ Q).squeeze(0) * Q).mean(1)
-
-        elif mode == 'mean-sign':
-            d = ((g.unsqueeze(0) @ Q).squeeze(0).sign() * Q).mean(1)
-
-        elif mode == 'mean-cosine':
-            d = (Q * _cosine_similarity(Q, g)).mean(1)
-
-        elif mode == 'mm':
-            d = (g.unsqueeze(0) @ Q).squeeze(0) / g.numel()
-
-        else:
-            raise ValueError(mode)
-
-        var.update = vec_to_tensors(g.dot(d).sign() * d, params)
-        return var
-

torchzero/modules/experimental/etf.py

@@ -1,172 +0,0 @@
-from typing import cast
-import warnings
-
-import torch
-
-from ...core import Module
-from ...utils import vec_to_tensors, vec_to_tensors_
-
-
-class ExponentialTrajectoryFit(Module):
-    """A method. Please note that this is experimental and isn't guaranteed to work."""
-    def __init__(self, step_size=1e-3):
-        defaults = dict(step_size = step_size)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        assert closure is not None
-        step_size = self.settings[var.params[0]]['step_size']
-
-        # 1. perform 3 GD steps to obtain 4 points
-        points = [torch.cat([p.view(-1) for p in var.params])]
-        for i in range(3):
-            if i == 0: grad = var.get_grad()
-            else:
-                with torch.enable_grad(): closure()
-                grad = [cast(torch.Tensor, p.grad) for p in var.params]
-
-            # GD step
-            torch._foreach_sub_(var.params, grad, alpha=step_size)
-
-            points.append(torch.cat([p.view(-1) for p in var.params]))
-
-        assert len(points) == 4, len(points)
-        x0, x1, x2, x3 = points
-        dim = x0.numel()
-
-        # 2. fit a generalized exponential curve
-        d0 = (x1 - x0).unsqueeze(1) # column vectors
-        d1 = (x2 - x1).unsqueeze(1)
-        d2 = (x3 - x2).unsqueeze(1)
-
-        # cat
-        D1 = torch.cat([d0, d1], dim=1)
-        D2 = torch.cat([d1, d2], dim=1)
-
-        # if points are collinear this will happen on sphere and a quadratic "line search" will minimize it
-        if x0.numel() >= 2:
-            if torch.linalg.matrix_rank(D1) < 2: # pylint:disable=not-callable
-                pass # need to put a quadratic fit there
-
-        M = D2 @ torch.linalg.pinv(D1) # pylint:disable=not-callable # this defines the curve
-
-        # now we can predict x*
-        I = torch.eye(dim, device=x0.device, dtype=x0.dtype)
-        B = I - M
-        z = x1 - M @ x0
-
-        x_star = torch.linalg.lstsq(B, z).solution # pylint:disable=not-callable
-
-        vec_to_tensors_(x0, var.params)
-        difference = torch._foreach_sub(var.params, vec_to_tensors(x_star, var.params))
-        var.update = list(difference)
-        return var
-
-
-
-class ExponentialTrajectoryFitV2(Module):
-    """Should be better than one above, except it isn't. Please note that this is experimental and isn't guaranteed to work."""
-    def __init__(self, step_size=1e-3, num_steps: int= 4):
-        defaults = dict(step_size = step_size, num_steps=num_steps)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        assert closure is not None
-        step_size = self.settings[var.params[0]]['step_size']
-        num_steps = self.settings[var.params[0]]['num_steps']
-
-        # 1. perform 3 GD steps to obtain 4 points (or more)
-        grad = var.get_grad()
-        points = [torch.cat([p.view(-1) for p in var.params])]
-        point_grads = [torch.cat([g.view(-1) for g in grad])]
-
-        for i in range(num_steps):
-            # GD step
-            torch._foreach_sub_(var.params, grad, alpha=step_size)
-
-            points.append(torch.cat([p.view(-1) for p in var.params]))
-
-            closure(backward=True)
-            grad = [cast(torch.Tensor, p.grad) for p in var.params]
-            point_grads.append(torch.cat([g.view(-1) for g in grad]))
-
-
-        X = torch.stack(points, 1) # dim, num_steps+1
-        G = torch.stack(point_grads, 1)
-        dim = points[0].numel()
-
-        X = torch.cat([X, torch.ones(1, num_steps+1, dtype=G.dtype, device=G.device)])
-
-        P = G @ torch.linalg.pinv(X) # pylint:disable=not-callable
-        A = P[:, :dim]
-        b = -P[:, dim]
-
-        # symmetrize
-        A = 0.5 * (A + A.T)
-
-        # predict x*
-        x_star = torch.linalg.lstsq(A, b).solution # pylint:disable=not-callable
-
-        vec_to_tensors_(points[0], var.params)
-        difference = torch._foreach_sub(var.params, vec_to_tensors(x_star, var.params))
-        var.update = list(difference)
-        return var
-
-
-
-
-def _fit_exponential(y0, y1, y2):
-    """x0, x1 and x2 are assumed to be 0, 1, 2"""
-    r = (y2 - y1) / (y1 - y0)
-    ones = r==1
-    r[ones] = 0
-    B = (y1 - y0) / (r - 1)
-    A = y0 - B
-
-    A[ones] = 0
-    B[ones] = 0
-    return A, B, r
-
-class PointwiseExponential(Module):
-    """A stupid method (for my youtube channel). Please note that this is experimental and isn't guaranteed to work."""
-    def __init__(self, step_size: float = 1e-3, reg: float = 1e-2, steps = 10000):
-        defaults = dict(reg=reg, steps=steps, step_size=step_size)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        assert closure is not None
-        settings = self.settings[var.params[0]]
-        step_size = settings['step_size']
-        reg = settings['reg']
-        steps = settings['steps']
-
-        # 1. perform 2 GD steps to obtain 3 points
-        points = [torch.cat([p.view(-1) for p in var.params])]
-        for i in range(2):
-            if i == 0: grad = var.get_grad()
-            else:
-                with torch.enable_grad(): closure()
-                grad = [cast(torch.Tensor, p.grad) for p in var.params]
-
-            # GD step
-            torch._foreach_sub_(var.params, grad, alpha=step_size)
-
-            points.append(torch.cat([p.view(-1) for p in var.params]))
-
-        assert len(points) == 3, len(points)
-        y0, y1, y2 = points
-
-        A, B, r = _fit_exponential(y0, y1, y2)
-        r = r.clip(max = 1-reg)
-        x_star = A + B * r**steps
-
-        vec_to_tensors_(y0, var.params)
-        difference = torch._foreach_sub(var.params, vec_to_tensors(x_star, var.params))
-        var.update = list(difference)
-        return var

torchzero/modules/experimental/soapy.py

@@ -1,163 +0,0 @@
-from operator import itemgetter
-
-import torch
-
-from ...core import Chainable, Transform
-from ..optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
-from ..optimizers.soap import (
-    update_soap_covariances_,
-    get_orthogonal_matrix,
-    get_orthogonal_matrix_QR,
-    project,
-    project_back,
-)
-
-class SOAPY(Transform):
-    """Adam but uses scaled gradient differences for GGᵀ. Please note that this is experimental and isn't guaranteed to work.
-
-    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
-    """
-    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,
-        y_to_ema2: 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,
-            y_to_ema2=y_to_ema2,
-        )
-        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']
-            y_to_ema2 = setting['y_to_ema2']
-
-            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
-            if scale_by_s: y /= torch.linalg.norm(s).clip(min=1e-8) # pylint:disable=not-callable
-
-            state['p_prev'].copy_(p)
-            state['g_prev'].copy_(t)
-
-            # initialize state on 1st step
-            if 'GG' not in state:
-                state["exp_avg"] = torch.zeros_like(t)
-                if y_to_ema2: state["exp_avg_sq"] = torch.ones_like(t)
-                else: 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_soap_covariances_(y, 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']
-            z_projected = None
-            if state['GG'] is not None:
-                if y_to_ema2: z_projected = project(y, state['Q'])
-                else: z_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"]
-
-            exp_avg.lerp_(t, 1-beta1)
-
-            if z_projected is None:
-                if y_to_ema2: exp_avg_sq.mul_(beta2).addcmul_(y, y, value=1-beta2)
-                else: exp_avg_sq.mul_(beta2).addcmul_(t, t, value=1-beta2)
-            else:
-                exp_avg_sq.mul_(beta2).addcmul_(z_projected, z_projected, value=1-beta2)
-
-            # project exponential moving averages if they are accumulated unprojected
-            exp_avg_projected = exp_avg
-            if z_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 z_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_soap_covariances_(y, 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