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

torchzero/modules/experimental/structured_newton.py

@@ -1,111 +0,0 @@
-# idea https://arxiv.org/pdf/2212.09841
-import warnings
-from collections.abc import Callable
-from functools import partial
-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,
-    hessian_mat,
-    hvp,
-    hvp_fd_central,
-    hvp_fd_forward,
-    jacobian_and_hessian_wrt,
-)
-
-
-class StructuredNewton(Module):
-    """TODO. Please note that this is experimental and isn't guaranteed to work.
-    Args:
-        structure (str, optional): structure.
-        reg (float, optional): tikhonov regularizer value. Defaults to 1e-6.
-        hvp_method (str):
-            how to calculate hvp_method. Defaults to "autograd".
-        inner (Chainable | None, optional): inner modules. Defaults to None.
-
-    """
-    def __init__(
-        self,
-        structure: Literal[
-            "diagonal",
-            "diagonal1",
-            "diagonal_abs",
-            "tridiagonal",
-            "circulant",
-            "toeplitz",
-            "toeplitz_like",
-            "hankel",
-            "rank1",
-            "rank2", # any rank
-        ]
-        | str = "diagonal",
-        reg: float = 1e-6,
-        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
-        h: float = 1e-3,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(reg=reg, hvp_method=hvp_method, structure=structure, h=h)
-        super().__init__(defaults)
-
-        if inner is not None:
-            self.set_child('inner', inner)
-
-    @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]]
-        reg = settings['reg']
-        hvp_method = settings['hvp_method']
-        structure = settings['structure']
-        h = settings['h']
-
-        # ------------------------ calculate grad and hessian ------------------------ #
-        if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
-            def Hvp_fn1(x):
-                return hvp(params, grad, x, retain_graph=True)
-            Hvp_fn = Hvp_fn1
-
-        elif hvp_method == 'forward':
-            grad = var.get_grad()
-            def Hvp_fn2(x):
-                return hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1]
-            Hvp_fn = Hvp_fn2
-
-        elif hvp_method == 'central':
-            grad = var.get_grad()
-            def Hvp_fn3(x):
-                return hvp_fd_central(closure, params, x, h=h, normalize=True)[1]
-            Hvp_fn = Hvp_fn3
-
-        else: raise ValueError(hvp_method)
-
-        # -------------------------------- inner step -------------------------------- #
-        update = var.get_update()
-        if 'inner' in self.children:
-            update = apply_transform(self.children['inner'], update, params=params, grads=grad, var=var)
-
-        # hessian
-        if structure.startswith('diagonal'):
-            H = Hvp_fn([torch.ones_like(p) for p in params])
-            if structure == 'diagonal1': torch._foreach_clamp_min_(H, 1)
-            if structure == 'diagonal_abs': torch._foreach_abs_(H)
-            torch._foreach_add_(H, reg)
-            torch._foreach_div_(update, H)
-            var.update = update
-            return var
-
-        # hessian
-        raise NotImplementedError(structure)
-
-
-
-
-

torchzero/modules/experimental/subspace_preconditioners.py

@@ -1,138 +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. Please note that this is experimental and isn't guaranteed to work."""
-    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(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.
-    Please note that this is experimental and isn't guaranteed to work.
-
-    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(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/tada.py

@@ -1,38 +0,0 @@
-from collections import deque
-
-import torch
-
-from ...core import Chainable, TensorwiseTransform
-from ...utils.linalg import matrix_power_eigh
-
-
-class TAda(TensorwiseTransform):
-    """3rd order whitening (maybe normalizes skewness). Please note that this is experimental and isn't guaranteed to work."""
-    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, settings):
-        reg = settings['reg']
-        if 'history' not in state:
-            state['history'] = deque(maxlen=settings['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, settings):
-        outer = state['outer']
-        P = matrix_power_eigh(outer, -1/2)
-        return (P @ tensor.ravel()).view_as(tensor)

torchzero/modules/line_search/trust_region.py

@@ -1,73 +0,0 @@
-from operator import itemgetter
-
-import torch
-
-from .line_search import LineSearch
-
-
-class TrustRegion(LineSearch):
-    """Basic first order trust region method. Re-evaluates the function after stepping, if value decreased sufficiently,
-    step size is increased. If value increased, step size is decreased. This is prone to collapsing.
-
-    Args:
-        nplus (float, optional): multiplier to step size on successful steps. Defaults to 1.5.
-        nminus (float, optional): multiplier to step size on unsuccessful steps. Defaults to 0.75.
-        c (float, optional): descent condition. Defaults to 1e-4.
-        init (float, optional): initial step size. Defaults to 1.
-        backtrack (bool, optional): whether to undo the step if value increased. Defaults to True.
-        adaptive (bool, optional):
-            If enabled, when multiple consecutive steps have been successful or unsuccessful,
-            the corresponding multipliers are increased, otherwise they are reset. Defaults to True.
-    """
-    def __init__(self, nplus: float=1.5, nminus: float=0.75, c: float=1e-4, init: float = 1, backtrack: bool = True, adaptive: bool = True):
-        defaults = dict(nplus=nplus, nminus=nminus, c=c, init=init, backtrack=backtrack, adaptive=adaptive)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def search(self, update, var):
-
-        nplus, nminus, c, init, backtrack, adaptive = itemgetter('nplus','nminus','c','init','backtrack', 'adaptive')(self.settings[var.params[0]])
-        step_size = self.global_state.setdefault('step_size', init)
-        previous_success = self.global_state.setdefault('previous_success', False)
-        nplus_mul =  self.global_state.setdefault('nplus_mul', 1)
-        nminus_mul = self.global_state.setdefault('nminus_mul', 1)
-
-
-        f_0 = self.evaluate_step_size(0, var, backward=False)
-
-        # directional derivative (0 if c = 0 because it is not needed)
-        if c == 0: d = 0
-        else: d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), update))
-
-        # test step size
-        sufficient_f = f_0 + c * step_size * min(d, 0) # pyright:ignore[reportArgumentType]
-
-        f_1 = self.evaluate_step_size(step_size, var, backward=False)
-
-        proposed = step_size
-
-        # very good step
-        if f_1 < sufficient_f:
-            self.global_state['step_size'] *= nplus * nplus_mul
-
-            # two very good steps in a row - increase nplus_mul
-            if adaptive:
-                if previous_success: self.global_state['nplus_mul'] *= nplus
-                else: self.global_state['nplus_mul'] = 1
-
-        # acceptable step step
-        #elif f_1 <= f_0: pass
-
-        # bad step
-        if f_1 >= f_0:
-            self.global_state['step_size'] *= nminus * nminus_mul
-
-            # two bad steps in a row - decrease nminus_mul
-            if adaptive:
-                if previous_success: self.global_state['nminus_mul'] *= nminus
-                else: self.global_state['nminus_mul'] = 1
-
-            if backtrack: proposed = 0
-            else: proposed *= nminus * nminus_mul
-
-        return proposed

torchzero/modules/lr/__init__.py

@@ -1,2 +0,0 @@
-from .lr import LR, StepSize, Warmup
-from .adaptive import PolyakStepSize, RandomStepSize

torchzero/modules/lr/adaptive.py

@@ -1,93 +0,0 @@
-"""Various step size strategies"""
-import random
-from typing import Any
-from operator import itemgetter
-import torch
-
-from ...core import Transform
-from ...utils import TensorList, NumberList, unpack_dicts
-
-
-class PolyakStepSize(Transform):
-    """Polyak's step-size method.
-
-    Args:
-        max (float | None, optional): maximum possible step size. Defaults to None.
-        min_obj_value (int, optional):
-            (estimated) minimal possible value of the objective function (lowest possible loss). Defaults to 0.
-        use_grad (bool, optional):
-            if True, uses dot product of update and gradient to compute the step size.
-            Otherwise, dot product of update with itself is used, which has no geometric meaning so it probably won't work well.
-            Defaults to True.
-        parameterwise (bool, optional):
-            if True, calculate Polyak step-size for each parameter separately,
-            if False calculate one global step size for all parameters. Defaults to False.
-        alpha (float, optional): multiplier to Polyak step-size. Defaults to 1.
-    """
-    def __init__(self, max: float | None = None, min_obj_value: float = 0, use_grad=True, parameterwise=False, alpha: float = 1):
-
-        defaults = dict(alpha=alpha, max=max, min_obj_value=min_obj_value, use_grad=use_grad, parameterwise=parameterwise)
-        super().__init__(defaults, uses_grad=use_grad)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        assert grads is not None
-        tensors = TensorList(tensors)
-        grads = TensorList(grads)
-        alpha = NumberList(s['alpha'] for s in settings)
-
-        parameterwise, use_grad, max, min_obj_value = itemgetter('parameterwise', 'use_grad', 'max', 'min_obj_value')(settings[0])
-
-        if use_grad: denom = tensors.dot(grads)
-        else: denom = tensors.dot(tensors)
-
-        if parameterwise:
-            polyak_step_size: TensorList | Any = (loss - min_obj_value) / denom.where(denom!=0, 1)
-            polyak_step_size = polyak_step_size.where(denom != 0, 0)
-            if max is not None: polyak_step_size = polyak_step_size.clamp_max(max)
-
-        else:
-            if denom.abs() <= torch.finfo(denom.dtype).eps: polyak_step_size = 0 # converged
-            else: polyak_step_size = (loss - min_obj_value) / denom
-
-            if max is not None:
-                if polyak_step_size > max: polyak_step_size = max
-
-        tensors.mul_(alpha * polyak_step_size)
-        return tensors
-
-
-class RandomStepSize(Transform):
-    """Uses random global or layer-wise step size from `low` to `high`.
-
-    Args:
-        low (float, optional): minimum learning rate. Defaults to 0.
-        high (float, optional): maximum learning rate. Defaults to 1.
-        parameterwise (bool, optional):
-            if True, generate random step size for each parameter separately,
-            if False generate one global random step size. Defaults to False.
-    """
-    def __init__(self, low: float = 0, high: float = 1, parameterwise=False, seed:int|None=None):
-        defaults = dict(low=low, high=high, parameterwise=parameterwise,seed=seed)
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        s = settings[0]
-        parameterwise = s['parameterwise']
-
-        seed = s['seed']
-        if 'generator' not in self.global_state:
-            self.global_state['generator'] = random.Random(seed)
-        generator: random.Random = self.global_state['generator']
-
-        if parameterwise:
-            low, high = unpack_dicts(settings, 'low', 'high')
-            lr = [generator.uniform(l, h) for l, h in zip(low, high)]
-        else:
-            low = s['low']
-            high = s['high']
-            lr = generator.uniform(low, high)
-
-        torch._foreach_mul_(tensors, lr)
-        return tensors

torchzero/modules/lr/lr.py

@@ -1,63 +0,0 @@
-"""Learning rate"""
-import torch
-
-from ...core import Transform
-from ...utils import NumberList, TensorList, generic_eq, unpack_dicts
-
-def lazy_lr(tensors: TensorList, lr: float | list, inplace:bool):
-    """multiplies by lr if lr is not 1"""
-    if generic_eq(lr, 1): return tensors
-    if inplace: return tensors.mul_(lr)
-    return tensors * lr
-
-class LR(Transform):
-    """Learning rate. Adding this module also adds support for LR schedulers."""
-    def __init__(self, lr: float):
-        defaults=dict(lr=lr)
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        return lazy_lr(TensorList(tensors), lr=[s['lr'] for s in settings], inplace=True)
-
-class StepSize(Transform):
-    """this is exactly the same as LR, except the `lr` parameter can be renamed to any other name to avoid clashes"""
-    def __init__(self, step_size: float, key = 'step_size'):
-        defaults={"key": key, key: step_size}
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        return lazy_lr(TensorList(tensors), lr=[s[s['key']] for s in settings], inplace=True)
-
-
-def _warmup_lr(step: int, start_lr: float | NumberList, end_lr: float | NumberList, steps: float):
-    """returns warm up lr scalar"""
-    if step > steps: return end_lr
-    return start_lr + (end_lr - start_lr) * (step / steps)
-
-class Warmup(Transform):
-    """Learning rate warmup, linearly increases learning rate multiplier from :code:`start_lr` to :code:`end_lr` over :code:`steps` steps.
-
-    Args:
-        start_lr (_type_, optional): initial learning rate multiplier on first step. Defaults to 1e-5.
-        end_lr (float, optional): learning rate multiplier at the end and after warmup. Defaults to 1.
-        steps (int, optional): number of steps to perform warmup for. Defaults to 100.
-    """
-    def __init__(self, start_lr = 1e-5, end_lr:float = 1, steps = 100):
-        defaults = dict(start_lr=start_lr,end_lr=end_lr, steps=steps)
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        start_lr, end_lr = unpack_dicts(settings, 'start_lr', 'end_lr', cls = NumberList)
-        num_steps = settings[0]['steps']
-        step = self.global_state.get('step', 0)
-
-        target = lazy_lr(
-            TensorList(tensors),
-            lr=_warmup_lr(step=step, start_lr=start_lr, end_lr=end_lr, steps=num_steps),
-            inplace=True
-        )
-        self.global_state['step'] = step + 1
-        return target

torchzero/modules/momentum/matrix_momentum.py

@@ -1,166 +0,0 @@
-from typing import Literal
-
-import torch
-
-from ...core import Module, apply_transform, Chainable
-from ...utils import NumberList, TensorList, as_tensorlist
-from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
-
-class MatrixMomentum(Module):
-    """
-    May be useful for ill conditioned stochastic quadratic objectives but I need to test this.
-    Evaluates hessian vector product on each step (via finite difference or autograd).
-
-    `mu` is supposed to be smaller than (1/largest eigenvalue), otherwise this will be very unstable.
-
-    Args:
-        mu (float, optional): this has a similar role to (1 - beta) in normal momentum. Defaults to 0.1.
-        beta (float, optional): decay for the buffer, this is not part of the original update rule. Defaults to 1.
-        hvp_method (str, optional):
-            How to calculate hessian-vector products.
-            Exact - "autograd", or finite difference - "forward", "central". Defaults to 'forward'.
-        h (float, optional): finite difference step size if hvp_method is set to finite difference. Defaults to 1e-3.
-        hvp_tfm (Chainable | None, optional): optional module applied to hessian-vector products. Defaults to None.
-
-    Reference:
-        Orr, Genevieve, and Todd Leen. "Using curvature information for fast stochastic search." Advances in neural information processing systems 9 (1996).
-    """
-
-    def __init__(
-        self,
-        mu=0.1,
-        beta: float = 1,
-        hvp_method: Literal["autograd", "forward", "central"] = "forward",
-        h: float = 1e-3,
-        hvp_tfm: Chainable | None = None,
-    ):
-        defaults = dict(mu=mu, beta=beta, hvp_method=hvp_method, h=h)
-        super().__init__(defaults)
-
-        if hvp_tfm is not None:
-            self.set_child('hvp_tfm', hvp_tfm)
-
-    @torch.no_grad
-    def step(self, var):
-        assert var.closure is not None
-        prev_update = self.get_state(var.params, 'prev_update', cls=TensorList)
-        hvp_method = self.settings[var.params[0]]['hvp_method']
-        h = self.settings[var.params[0]]['h']
-
-        mu,beta = self.get_settings(var.params, 'mu','beta', cls=NumberList)
-
-        if hvp_method == 'autograd':
-            with torch.enable_grad():
-                grad = var.get_grad(create_graph=True)
-                hvp_ = TensorList(hvp(var.params, grads=grad, vec=prev_update, allow_unused=True, retain_graph=False)).detach_()
-
-        elif hvp_method == 'forward':
-            var.get_grad()
-            l, hvp_ = hvp_fd_forward(var.closure, var.params, vec=prev_update, g_0=var.grad, h=h, normalize=True)
-            if var.loss_approx is None: var.loss_approx = l
-
-        elif hvp_method == 'central':
-            l, hvp_ = hvp_fd_central(var.closure, var.params, vec=prev_update, h=h, normalize=True)
-            if var.loss_approx is None: var.loss_approx = l
-
-        else:
-            raise ValueError(hvp_method)
-
-        if 'hvp_tfm' in self.children:
-            hvp_ = TensorList(apply_transform(self.children['hvp_tfm'], hvp_, params=var.params, grads=var.grad, var=var))
-
-        update = TensorList(var.get_update())
-
-        hvp_ = as_tensorlist(hvp_)
-        update.add_(prev_update - hvp_*mu)
-        prev_update.set_(update * beta)
-        var.update = update
-        return var
-
-
-class AdaptiveMatrixMomentum(Module):
-    """
-    May be useful for ill conditioned stochastic quadratic objectives but I need to test this.
-    Evaluates hessian vector product on each step (via finite difference or autograd).
-
-    This version estimates mu via a simple heuristic: ||s||/||y||, where s is parameter difference, y is gradient difference.
-
-    Args:
-        mu_mul (float, optional): multiplier to the estimated mu. Defaults to 1.
-        beta (float, optional): decay for the buffer, this is not part of the original update rule. Defaults to 1.
-        hvp_method (str, optional):
-            How to calculate hessian-vector products.
-            Exact - "autograd", or finite difference - "forward", "central". Defaults to 'forward'.
-        h (float, optional): finite difference step size if hvp_method is set to finite difference. Defaults to 1e-3.
-        hvp_tfm (Chainable | None, optional): optional module applied to hessian-vector products. Defaults to None.
-
-    Reference:
-        Orr, Genevieve, and Todd Leen. "Using curvature information for fast stochastic search." Advances in neural information processing systems 9 (1996).
-    """
-
-    def __init__(
-        self,
-        mu_mul: float = 1,
-        beta: float = 1,
-        eps=1e-4,
-        hvp_method: Literal["autograd", "forward", "central"] = "forward",
-        h: float = 1e-3,
-        hvp_tfm: Chainable | None = None,
-    ):
-        defaults = dict(mu_mul=mu_mul, beta=beta, hvp_method=hvp_method, h=h, eps=eps)
-        super().__init__(defaults)
-
-        if hvp_tfm is not None:
-            self.set_child('hvp_tfm', hvp_tfm)
-
-    @torch.no_grad
-    def step(self, var):
-        assert var.closure is not None
-        prev_update, prev_params, prev_grad = self.get_state(var.params, 'prev_update', 'prev_params', 'prev_grad', cls=TensorList)
-
-        settings = self.settings[var.params[0]]
-        hvp_method = settings['hvp_method']
-        h = settings['h']
-        eps = settings['eps']
-
-        mu_mul, beta = self.get_settings(var.params, 'mu_mul','beta', cls=NumberList)
-
-        if hvp_method == 'autograd':
-            with torch.enable_grad():
-                grad = var.get_grad(create_graph=True)
-                hvp_ = TensorList(hvp(var.params, grads=grad, vec=prev_update, allow_unused=True, retain_graph=False)).detach_()
-
-        elif hvp_method == 'forward':
-            var.get_grad()
-            l, hvp_ = hvp_fd_forward(var.closure, var.params, vec=prev_update, g_0=var.grad, h=h, normalize=True)
-            if var.loss_approx is None: var.loss_approx = l
-
-        elif hvp_method == 'central':
-            l, hvp_ = hvp_fd_central(var.closure, var.params, vec=prev_update, h=h, normalize=True)
-            if var.loss_approx is None: var.loss_approx = l
-
-        else:
-            raise ValueError(hvp_method)
-
-        if 'hvp_tfm' in self.children:
-            hvp_ = TensorList(apply_transform(self.children['hvp_tfm'], hvp_, params=var.params, grads=var.grad, var=var))
-
-        # adaptive part
-        update = TensorList(var.get_update())
-
-        s_k = var.params - prev_params
-        prev_params.copy_(var.params)
-
-        assert var.grad is not None
-        y_k = var.grad - prev_grad
-        prev_grad.copy_(var.grad)
-
-        ada_mu = (s_k.global_vector_norm() / (y_k.global_vector_norm() + eps)) * mu_mul
-
-        # matrix momentum uppdate
-        hvp_ = as_tensorlist(hvp_)
-        update.add_(prev_update - hvp_*ada_mu)
-        prev_update.set_(update * beta)
-        var.update = update
-        return var
-