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

torchzero/modules/experimental/modular_lbfgs.py

@@ -1,265 +0,0 @@
-from collections import deque
-from operator import itemgetter
-from typing import Any
-
-import torch
-
-from ...core import Chainable, Module, Transform, Var, apply_transform, maybe_chain
-from ...utils import NumberList, TensorList, as_tensorlist
-
-
-def _adaptive_damping(
-    s_k: TensorList,
-    y_k: TensorList,
-    ys_k: torch.Tensor,
-    init_damping = 0.99,
-    eigval_bounds = (0.01, 1.5)
-):
-    # adaptive damping Al-Baali, M.: Quasi-Wolfe conditions for quasi-Newton methods for large-scale optimization. In: 40th Workshop on Large Scale Nonlinear Optimization, Erice, Italy, June 22–July 1 (2004)
-    sigma_l, sigma_h = eigval_bounds
-    u = ys_k / s_k.dot(s_k)
-    if u <= sigma_l < 1: tau = min((1-sigma_l)/(1-u), init_damping)
-    elif u >= sigma_h > 1: tau = min((sigma_h-1)/(u-1), init_damping)
-    else: tau = init_damping
-    y_k = tau * y_k + (1-tau) * s_k
-    ys_k = s_k.dot(y_k)
-
-    return s_k, y_k, ys_k
-
-def lbfgs(
-    tensors_: TensorList,
-    var: Var,
-    s_history: deque[TensorList],
-    y_history: deque[TensorList],
-    sy_history: deque[torch.Tensor],
-    y_k: TensorList | None,
-    ys_k: torch.Tensor | None,
-    z_tfm: Any,
-):
-    if len(s_history) == 0 or y_k is None or ys_k is None:
-
-        # initial step size guess modified from pytorch L-BFGS
-        scale = 1 / tensors_.abs().global_sum()
-        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
-        return tensors_.mul_(min(1.0, scale)) # pyright: ignore[reportArgumentType]
-
-    # 1st loop
-    alpha_list = []
-    q = tensors_.clone()
-    for s_i, y_i, ys_i in zip(reversed(s_history), reversed(y_history), reversed(sy_history)):
-        p_i = 1 / ys_i # this is also denoted as ρ (rho)
-        alpha = p_i * s_i.dot(q)
-        alpha_list.append(alpha)
-        q.sub_(y_i, alpha=alpha) # pyright: ignore[reportArgumentType]
-
-    # calculate z
-    # s.y/y.y is also this weird y-looking symbol I couldn't find
-    # z is it times q
-    # actually H0 = (s.y/y.y) * I, and z = H0 @ q
-    z = q * (ys_k / (y_k.dot(y_k)))
-
-    if z_tfm is not None:
-        z = TensorList(apply_transform(z_tfm, tensors=z, params=var.params, grads=var.grad, var=var))
-
-    # 2nd loop
-    for s_i, y_i, ys_i, alpha_i in zip(s_history, y_history, sy_history, reversed(alpha_list)):
-        p_i = 1 / ys_i
-        beta_i = p_i * y_i.dot(z)
-        z.add_(s_i, alpha = alpha_i - beta_i)
-
-    return z
-
-def _apply_tfms_into_history(
-    self: Module,
-    params: list[torch.Tensor],
-    var: Var,
-    update: list[torch.Tensor],
-):
-    if 'params_history_tfm' in self.children:
-        params = apply_transform(self.children['params_history_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=var.grad, var=var)
-
-    if 'grad_history_tfm' in self.children:
-        update = apply_transform(self.children['grad_history_tfm'], tensors=as_tensorlist(update).clone(), params=params, grads=var.grad, var=var)
-
-    return params, update
-
-def _apply_tfms_into_precond(
-    self: Module,
-    params: list[torch.Tensor],
-    var: Var,
-    update: list[torch.Tensor],
-):
-    if 'params_precond_tfm' in self.children:
-        params = apply_transform(self.children['params_precond_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=var.grad, var=var)
-
-    if 'grad_precond_tfm' in self.children:
-        update = apply_transform(self.children['grad_precond_tfm'], tensors=update, params=params, grads=var.grad, var=var)
-
-    return params, update
-
-
-class ModularLBFGS(Module):
-    """L-BFGS with ability to apply transforms to many inner variables.
-
-    Args:
-        history_size (int, optional): number of past parameter differences and gradient differences to store. Defaults to 10.
-        tol (float | None, optional):
-            tolerance for minimal gradient difference to avoid instability after converging to minima. Defaults to 1e-10.
-        damping (bool, optional):
-            whether to use adaptive damping. Learning rate might need to be lowered with this enabled. Defaults to False.
-        init_damping (float, optional):
-            initial damping for adaptive dampening. Defaults to 0.9.
-        eigval_bounds (tuple, optional):
-            eigenvalue bounds for adaptive dampening. Defaults to (0.5, 50).
-        update_freq (int, optional):
-            how often to update L-BFGS history. Defaults to 1.
-        z_tfm (float | None, optional):
-            transform module applied to initial H^-1 @ q guess. Defaults to None.
-        params_history_tfm (AnyTransform | None, optional):
-            transform module applied to params before adding s_k to history. Defaults to None.
-        grad_history_tfm (AnyTransform | None, optional):
-            transform module applied to grads before adding y_k to history. Defaults to None.
-        params_precond_tfm (AnyTransform | None, optional):
-            transform module applied to params to calculate s_k before preconditioning. Defaults to None.
-        grad_precond_tfm (AnyTransform | None, optional):
-            transform module applied to grads to calculate y_k before preconditioning. Defaults to None.
-        update_precond_tfm (Chainable | None, optional):
-            transform module applied to grads that are being preconditioned. Defaults to None.
-    """
-    def __init__(
-        self,
-        history_size=10,
-        tol: float | None = 1e-10,
-        damping: bool = False,
-        init_damping=0.9,
-        eigval_bounds=(0.5, 50),
-        update_freq = 1,
-        params_history_tfm: Chainable | None = None,
-        grad_history_tfm: Chainable | None = None,
-        params_precond_tfm: Chainable | None = None,
-        grad_precond_tfm: Chainable | None = None,
-        update_precond_tfm: Chainable | None = None,
-        z_tfm: Chainable | None = None,
-    ):
-        defaults = dict(history_size=history_size, tol=tol, damping=damping, init_damping=init_damping, eigval_bounds=eigval_bounds, update_freq=update_freq)
-        super().__init__(defaults)
-
-        self.global_state['s_history'] = deque(maxlen=history_size)
-        self.global_state['y_history'] = deque(maxlen=history_size)
-        self.global_state['sy_history'] = deque(maxlen=history_size)
-
-        loc = locals().copy()
-        for k in ('update_precond_tfm', 'params_history_tfm', 'grad_history_tfm', 'params_precond_tfm', 'grad_precond_tfm','z_tfm'):
-            v = loc[k]
-            if v is not None:
-                self.set_child(k,v)
-
-    def reset(self):
-        """Resets the internal state of the L-SR1 module."""
-        # super().reset() # Clears self.state (per-parameter) if any, and "step"
-        self.state.clear()
-        self.global_state['step'] = 0
-        self.global_state['s_history'].clear()
-        self.global_state['y_history'].clear()
-        self.global_state['sy_history'].clear()
-
-    @torch.no_grad
-    def step(self, var):
-        params = as_tensorlist(var.params)
-        update = as_tensorlist(var.get_update())
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
-
-        # history of s and k
-        s_history: deque[TensorList] = self.global_state['s_history']
-        y_history: deque[TensorList] = self.global_state['y_history']
-        sy_history: deque[torch.Tensor] = self.global_state['sy_history']
-
-        tol, damping, init_damping, eigval_bounds, update_freq = itemgetter(
-            'tol', 'damping', 'init_damping', 'eigval_bounds', 'update_freq')(self.settings[params[0]])
-
-        # params_beta, grads_beta = self.get_settings('params_beta', 'grads_beta', params=params, cls=NumberList)
-        # l_params, l_update = _lerp_params_update_(self, params, update, params_beta, grads_beta)
-
-        # params and update that go into history
-        params_h, update_h = _apply_tfms_into_history(
-            self,
-            params=params,
-            var=var,
-            update=update,
-        )
-
-        prev_params_h, prev_grad_h = self.get_state(params, 'prev_params_h', 'prev_grad_h', cls=TensorList)
-
-        # 1st step - there are no previous params and grads, `lbfgs` will do normalized SGD step
-        if step == 0:
-            s_k_h = None; y_k_h = None; ys_k_h = None
-        else:
-            s_k_h = params_h - prev_params_h
-            y_k_h = update_h - prev_grad_h
-            ys_k_h = s_k_h.dot(y_k_h)
-
-            if damping:
-                s_k_h, y_k_h, ys_k_h = _adaptive_damping(s_k_h, y_k_h, ys_k_h, init_damping=init_damping, eigval_bounds=eigval_bounds)
-
-        prev_params_h.copy_(params_h)
-        prev_grad_h.copy_(update_h)
-
-        # update effective preconditioning state
-        if step % update_freq == 0:
-            if ys_k_h is not None and ys_k_h > 1e-10:
-                assert s_k_h is not None and y_k_h is not None
-                s_history.append(s_k_h)
-                y_history.append(y_k_h)
-                sy_history.append(ys_k_h)
-
-        # step with inner module before applying preconditioner
-        if 'update_precond_tfm' in self.children:
-            update_precond_tfm = self.children['update_precond_tfm']
-            inner_var = update_precond_tfm.step(var.clone(clone_update=True))
-            var.update_attrs_from_clone_(inner_var)
-            tensors = inner_var.update
-            assert tensors is not None
-        else:
-            tensors = update.clone()
-
-        # transforms into preconditioner
-        params_p, update_p = _apply_tfms_into_precond(self, params=params, var=var, update=update)
-        prev_params_p, prev_grad_p = self.get_state(params, 'prev_params_p', 'prev_grad_p', cls=TensorList)
-
-        if step == 0:
-            s_k_p = None; y_k_p = None; ys_k_p = None
-
-        else:
-            s_k_p = params_p - prev_params_p
-            y_k_p = update_p - prev_grad_p
-            ys_k_p = s_k_p.dot(y_k_p)
-
-            if damping:
-                s_k_p, y_k_p, ys_k_p = _adaptive_damping(s_k_p, y_k_p, ys_k_p, init_damping=init_damping, eigval_bounds=eigval_bounds)
-
-        prev_params_p.copy_(params_p)
-        prev_grad_p.copy_(update_p)
-
-        # tolerance on gradient difference to avoid exploding after converging
-        if tol is not None:
-            if y_k_p is not None and y_k_p.abs().global_max() <= tol:
-                var.update = update # may have been updated by inner module, probably makes sense to use it here?
-                return var
-
-        # precondition
-        dir = lbfgs(
-            tensors_=as_tensorlist(tensors),
-            var=var,
-            s_history=s_history,
-            y_history=y_history,
-            sy_history=sy_history,
-            y_k=y_k_p,
-            ys_k=ys_k_p,
-            z_tfm=self.children.get('z_tfm', None),
-        )
-
-        var.update = dir
-
-        return var
-

torchzero/modules/experimental/parabolic_search.py

@@ -1,220 +0,0 @@
-import math
-from collections.abc import Mapping
-from operator import itemgetter
-
-import torch
-
-from ...core import Module
-from ...utils import TensorList
-
-
-
-def adaptive_tracking(
-    f,
-    f_0,
-    f_1,
-    t_0,
-    maxiter: int
-):
-
-    t = t_0
-    f_t = f(t)
-
-    # backtrack
-    if f_t > f_0:
-        if f_1 > f_0: t = min(0.5, t_0/2)
-        while f_t > f_0:
-            maxiter -= 1
-            if maxiter < 0: return 0, f_0
-            t = t/2
-            f_t = f(t) if t!=1 else f_1
-        return t, f_t
-
-    # forwardtrack
-    f_prev = f_t
-    t *= 2
-    f_t = f(t)
-    if f_prev < f_t: return t/2, f_prev
-    while f_prev >= f_t:
-        maxiter -= 1
-        if maxiter < 0: return t, f_t
-        f_prev = f_t
-        t *= 2
-        f_t = f(t)
-    return t/2, f_prev
-
-
-
-class ParabolaSearch(Module):
-    """"""
-    def __init__(
-        self,
-        step_size: float = 1e-2,
-        adaptive: bool=True,
-        normalize: bool=False,
-        # method: str | None = None,
-        maxiter: int | None = 10,
-        # bracket=None,
-        # bounds=None,
-        # tol: float | None = None,
-        # options=None,
-    ):
-        if normalize and adaptive: raise ValueError("pick either normalize or adaptive")
-        defaults = dict(step_size=step_size, adaptive=adaptive, normalize=normalize, maxiter=maxiter)
-        super().__init__(defaults)
-
-        import scipy.optimize
-        self.scopt = scipy.optimize
-
-
-    @torch.no_grad
-    def step(self, var):
-        x_0 = TensorList(var.params)
-        closure = var.closure
-        assert closure is not None
-        settings = self.settings[x_0[0]]
-        step_size = settings['step_size']
-        adaptive = settings['adaptive']
-        normalize = settings['normalize']
-        maxiter = settings['maxiter']
-        if normalize and adaptive: raise ValueError("pick either normalize or adaptive")
-
-        grad = TensorList(var.get_grad())
-        f_0 = var.get_loss(False)
-
-        scale = 1
-        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
-        if adaptive: scale = grad.abs().mean().clip(min=1e-8)
-
-        # make step
-        v_0 = grad * (step_size/scale)
-        x_0 -= v_0
-        with torch.enable_grad():
-            f_1 = closure()
-            grad = x_0.grad
-
-        x_0 += v_0
-        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
-        v_1 = grad * (step_size/scale)
-        a = v_1 - v_0
-
-        def parabolic_objective(t: float):
-            nonlocal x_0
-
-            step = v_0*t + 0.5*a*t**2
-            x_0 -= step
-            value = closure(False)
-            x_0 += step
-            return value.detach().cpu()
-
-        prev_t = self.global_state.get('prev_t', 2)
-        t, f = adaptive_tracking(parabolic_objective, f_0=f_0, f_1=f_1, t_0=prev_t, maxiter=maxiter)
-        self.global_state['prev_t'] = t
-
-        # method, bracket, bounds, tol, options, maxiter = itemgetter(
-        #     'method', 'bracket', 'bounds', 'tol', 'options', 'maxiter')(self.settings[var.params[0]])
-
-        # if maxiter is not None:
-        #     options = dict(options) if isinstance(options, Mapping) else {}
-        #     options['maxiter'] = maxiter
-
-        # res = self.scopt.minimize_scalar(parabolic_objective, method=method, bracket=bracket, bounds=bounds, tol=tol, options=options)
-        # t = res.x
-
-        var.update = v_0*t + 0.5*a*t**2
-        return var
-
-class CubicParabolaSearch(Module):
-    """"""
-    def __init__(
-        self,
-        step_size: float = 1e-2,
-        adaptive: bool=True,
-        normalize: bool=False,
-        # method: str | None = None,
-        maxiter: int | None = 10,
-        # bracket=None,
-        # bounds=None,
-        # tol: float | None = None,
-        # options=None,
-    ):
-        if normalize and adaptive: raise ValueError("pick either normalize or adaptive")
-        defaults = dict(step_size=step_size, adaptive=adaptive, normalize=normalize, maxiter=maxiter)
-        super().__init__(defaults)
-
-        import scipy.optimize
-        self.scopt = scipy.optimize
-
-
-    @torch.no_grad
-    def step(self, var):
-        x_0 = TensorList(var.params)
-        closure = var.closure
-        assert closure is not None
-        settings = self.settings[x_0[0]]
-        step_size = settings['step_size']
-        adaptive = settings['adaptive']
-        maxiter = settings['maxiter']
-        normalize = settings['normalize']
-        if normalize and adaptive: raise ValueError("pick either normalize or adaptive")
-
-        grad = TensorList(var.get_grad())
-        f_0 = var.get_loss(False)
-
-        scale = 1
-        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
-        if adaptive: scale = grad.abs().mean().clip(min=1e-8)
-
-        # make step
-        v_0 = grad * (step_size/scale)
-        x_0 -= v_0
-        with torch.enable_grad():
-            f_1 = closure()
-            grad = x_0.grad
-
-        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
-        v_1 = grad * (step_size/scale)
-        a_0 = v_1 - v_0
-
-        # make another step
-        x_0 -= v_1
-        with torch.enable_grad():
-            f_2 = closure()
-            grad = x_0.grad
-
-        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
-        v_2 = grad * (step_size/scale)
-        a_1 = v_2 - v_1
-
-        j = a_1 - a_0
-
-        x_0 += v_0
-        x_0 += v_1
-
-        def parabolic_objective(t: float):
-            nonlocal x_0
-
-            step = v_0*t + (1/2)*a_0*t**2 + (1/6)*j*t**3
-            x_0 -= step
-            value = closure(False)
-            x_0 += step
-            return value
-
-
-        prev_t = self.global_state.get('prev_t', 2)
-        t, f = adaptive_tracking(parabolic_objective, f_0=f_0, f_1=f_1, t_0=prev_t, maxiter=maxiter)
-        self.global_state['prev_t'] = t
-
-        # method, bracket, bounds, tol, options, maxiter = itemgetter(
-        #     'method', 'bracket', 'bounds', 'tol', 'options', 'maxiter')(self.settings[var.params[0]])
-
-        # if maxiter is not None:
-        #     options = dict(options) if isinstance(options, Mapping) else {}
-        #     options['maxiter'] = maxiter
-
-        # res = self.scopt.minimize_scalar(parabolic_objective, method=method, bracket=bracket, bounds=bounds, tol=tol, options=options)
-        # t = res.x
-
-        var.update = v_0*t + (1/2)*a_0*t**2 + (1/6)*j*t**3
-        return var
-

torchzero/modules/experimental/subspace_preconditioners.py

@@ -1,145 +0,0 @@
-from collections import deque
-
-import torch
-# import visualbench as vb
-
-# import torchzero as tz
-
-from ...core import Transform, Chainable, apply_transform
-from ...utils.linalg import inv_sqrt_2x2, matrix_power_eigh, gram_schmidt
-from ...utils import TensorList, vec_to_tensors_
-
-
-def inverse_sqrt(M):
-    if M.shape[-1] == 2: return inv_sqrt_2x2(M, force_pd=True) # general formula for 2x2 matrices
-    return matrix_power_eigh(M, -1/2)
-
-def update_subspace_preconditioner_(
-    grad: torch.Tensor, # store grads and basis as vectors for matmul
-    basis: torch.Tensor, # ndim, k
-    accumulator_: torch.Tensor, # k, k
-    beta: float | None,
-):
-    projected = basis.T @ grad # k
-    outer = torch.outer(projected, projected)
-
-    if beta is None: accumulator_.add_(outer)
-    else: accumulator_.lerp_(outer, 1-beta)
-
-def apply_subspace_preconditioner(
-    tensor: torch.Tensor,
-    basis: torch.Tensor, # ndim, k
-    accumulator: torch.Tensor,
-):
-    preconditioner = inverse_sqrt(accumulator) # k,k
-
-    tensor_projected = basis.T @ tensor # k
-    update_projected = preconditioner @ tensor_projected # k
-    return basis @ update_projected # d
-
-class RandomSubspacePreconditioning(Transform):
-    """Whitens in random slowly changing subspace.
-
-    .. warning::
-        Experimental and this is a barebones implementation.
-
-    """
-    def __init__(self, k: int, beta: float | None = 0.99, basis_beta: float | None = 0.99, inner: Chainable | None = None):
-        defaults = dict(k=k, beta=beta, basis_beta=basis_beta)
-        super().__init__(defaults, uses_grad=False)
-
-        if inner is not None: self.set_child('inner', inner)
-
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
-        settings = settings[0]
-        g = torch.cat([t.view(-1) for t in tensors])
-        k = settings['k']
-        beta = settings['beta']
-        basis_beta = settings['basis_beta']
-
-        if 'basis' not in self.global_state:
-            self.global_state['basis'] = torch.randn(g.numel(), k, device=g.device, dtype=g.dtype)
-            self.global_state['accumulator'] = torch.eye(k, device=g.device, dtype=g.dtype)
-
-        basis = self.global_state['basis']
-        accumulator = self.global_state['accumulator']
-
-        if basis_beta is not None:
-            basis.lerp_(torch.randn_like(basis), 1-basis_beta)
-
-        update_subspace_preconditioner_(g, basis, accumulator, beta)
-
-        if 'inner' in self.children:
-            tensors = apply_transform(self.children['inner'], tensors, params, grads)
-            g = torch.cat([t.view(-1) for t in tensors])
-
-        try:
-            preconditioned = apply_subspace_preconditioner(g, basis, accumulator)
-        except torch.linalg.LinAlgError:
-            preconditioned = g.clip(-0.1, 0.1)
-        vec_to_tensors_(preconditioned, tensors)
-
-        return tensors
-
-
-class HistorySubspacePreconditioning(Transform):
-    """Whitens in subspace spanned by history of gradient differences.
-
-    .. warning::
-        Experimental and this is a barebones implementation.
-
-    Args:
-        beta - for preconditioner itself in the basis.
-        basis_beta - how much basis is allowed to change.
-    """
-    def __init__(self, k: int, beta: float | None = 0.99, basis_beta=0.99, inner: Chainable | None = None):
-        defaults = dict(k=k, beta=beta, basis_beta=basis_beta)
-        super().__init__(defaults, uses_grad=False)
-
-        if inner is not None: self.set_child('inner', inner)
-
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
-        settings = settings[0]
-
-        g = torch.cat([t.view(-1) for t in tensors])
-        k = settings['k']
-        beta = settings['beta']
-        basis_beta = settings['basis_beta']
-
-        if 'history' not in self.global_state:
-            self.global_state['history'] = deque(maxlen=k)
-            self.global_state['accumulator'] = torch.eye(k, device=g.device, dtype=g.dtype)
-            self.global_state['basis'] = torch.ones(g.numel(), k, device=g.device, dtype=g.dtype)
-
-
-        history: deque = self.global_state['history']
-        accumulator = self.global_state['accumulator']
-        basis = self.global_state['basis']
-
-        history.append(g)
-        if len(history) < k:
-            basis_t = torch.randn(g.numel(), k, device=g.device, dtype=g.dtype)
-            history_basis = torch.stack(tuple(history), -1)
-            basis_t[:, -len(history):] = history_basis
-
-        else:
-            basis_t = torch.stack(tuple(history), -1)
-
-        basis_t[:,:-1] = basis_t[:, :-1] - basis_t[:, 1:]
-        basis_t = (basis_t - basis_t.mean()) / basis_t.std()
-
-        basis.lerp_(basis_t, 1-basis_beta)
-        update_subspace_preconditioner_(g, basis, accumulator, beta)
-
-        if 'inner' in self.children:
-            tensors = apply_transform(self.children['inner'], tensors, params, grads)
-            g = torch.cat([t.view(-1) for t in tensors])
-
-        try:
-            preconditioned = apply_subspace_preconditioner(g, basis, accumulator)
-        except torch.linalg.LinAlgError:
-            preconditioned = g.clip(-0.1,0.1)
-        vec_to_tensors_(preconditioned, tensors)
-
-        return tensors
-

torchzero/modules/experimental/tensor_adagrad.py

@@ -1,42 +0,0 @@
-from collections import deque
-
-import torch
-
-from ...core import Chainable, TensorwiseTransform
-from ...utils.linalg import matrix_power_eigh
-
-
-class TensorAdagrad(TensorwiseTransform):
-    """3rd order whitening (maybe normalizes skewness, but don't quote me on it).
-
-    .. warning::
-        Experimental.
-    """
-    def __init__(self, history_size: int = 100, reg: float = 1e-8, update_freq: int = 1, concat_params: bool = True, inner: Chainable | None = None):
-        defaults = dict(history_size=history_size, reg=reg)
-        super().__init__(defaults, uses_grad=False, update_freq=update_freq, inner=inner, concat_params=concat_params)
-
-    @torch.no_grad
-    def update_tensor(self, tensor, param, grad, loss, state, setting):
-        reg = setting['reg']
-        if 'history' not in state:
-            state['history'] = deque(maxlen=setting['history_size'])
-
-        g = tensor.view(-1)
-        history = state['history']
-        history.append(g.clone())
-
-        I = torch.eye(tensor.numel(), device=tensor.device, dtype=tensor.dtype).mul_(reg)
-        g_k = history[0]
-        outer = torch.outer(g_k, g_k).mul_(torch.dot(g_k, g).clip(min=reg))
-        if len(history) > 1:
-            for g_k in list(history)[1:]:
-                outer += torch.outer(g_k, g_k).mul_(torch.dot(g_k, g).clip(min=reg))
-
-        state['outer'] = outer.add_(I)
-
-    @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, setting):
-        outer = state['outer']
-        P = matrix_power_eigh(outer, -1/2)
-        return (P @ tensor.ravel()).view_as(tensor)