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

torchzero/modules/optimizers/adahessian.py

@@ -0,0 +1,223 @@
+import math
+from collections.abc import Callable
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, Target, Transform, apply_transform
+from ...utils import NumberList, TensorList, as_tensorlist
+from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
+
+
+def _block_average(x: torch.Tensor, block_size: int | None, enable: bool):
+    """averages x over first dimension in blocks"""
+    if enable and x.ndim >= 2:
+        if math.prod(x.shape[1:]) <= 1: return x
+        size = x.size(0)
+        if block_size is None: return x.mean(0, keepdim=True)
+
+        n_blocks = size // block_size
+        if n_blocks <= 1: return x.mean(0, keepdim = True)
+
+        n_remaining = size - n_blocks * block_size
+        remaining = None
+        if n_remaining > 0:
+            remaining = x[-n_remaining:].mean(0, keepdim=True).repeat_interleave(n_remaining, 0)
+            x = x[:-n_remaining]
+
+        x = x.view(block_size, n_blocks, *x.shape[1:])
+        x_mean = x.mean(0).repeat_interleave(block_size, 0)
+
+        if remaining is None: return x_mean
+        return torch.cat([x_mean, remaining], 0)
+
+    return x
+
+def _rademacher_like(tensor, p = 0.5, generator = None):
+    """p is probability of a 1, other values will be -1."""
+    return torch.bernoulli(torch.full_like(tensor, p), generator = generator).mul_(2).sub_(1)
+
+def adahessian(
+    tensors: TensorList,
+    D: TensorList | None,
+    exp_avg_: TensorList,
+    D_exp_avg_sq_: TensorList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    update_freq: int,
+    eps: float | NumberList,
+    step: int,
+):
+    # momentum
+    exp_avg_.lerp_(tensors, 1-beta1)
+    num = exp_avg_ / (1-beta1)
+
+    # update preconditioner
+    if step % update_freq == 0:
+        assert D is not None
+        D_exp_avg_sq_.mul_(beta2).addcmul_(D, D, 1-beta2)
+
+    else:
+        assert D is None
+
+    denom = (D_exp_avg_sq_ / (1-beta2)).sqrt_().add_(eps)
+
+    return num.div_(denom)
+
+
+class AdaHessian(Module):
+    """AdaHessian: An Adaptive Second Order Optimizer for Machine Learning (https://arxiv.org/abs/2006.00719)
+
+    This is similar to Adam, but the second momentum is replaced by square root of an exponential moving average of random hessian-vector products.
+
+    .. note::
+        In most cases AdaHessian should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply AdaHessian preconditioning to another module's output.
+
+    .. note::
+        If you are using gradient estimators or reformulations, set :code:`hvp_method` to "forward" or "central".
+
+    .. note::
+        This module requires a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients for calculating HVPs.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    Args:
+        beta1 (float, optional): first momentum. Defaults to 0.9.
+        beta2 (float, optional): second momentum for squared hessian diagonal estimates. Defaults to 0.999.
+        averaging (bool, optional):
+            whether to enable block diagonal averaging over 1st dimension on parameters that have 2+ dimensions.
+            This can be set per-parameter in param groups.
+        block_size (int, optional):
+            size of block in the block-diagonal averaging.
+        update_freq (int, optional):
+            frequency of updating hessian diagonal estimate via a hessian-vector product.
+            This value can be increased to reduce computational cost. Defaults to 1.
+        eps (float, optional):
+            division stability epsilon. Defaults to 1e-8.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        h (float, optional): finite difference step size if :code:`hvp_method` is "forward" or "central". Defaults to 1e-3.
+        n_samples (int, optional):
+            number of hessian-vector products with random vectors to evaluate each time when updating
+            the preconditioner. Larger values may lead to better hessian diagonal estimate. Defaults to 1.
+        seed (int | None, optional): seed for random vectors. Defaults to None.
+        inner (Chainable | None, optional):
+            Inner module. If this is specified, operations are performed in the following order.
+            1. compute hessian diagonal estimate.
+            2. pass inputs to :code:`inner`.
+            3. momentum and preconditioning are applied to the ouputs of :code:`inner`.
+
+    Examples:
+        Using AdaHessian:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.AdaHessian(),
+                tz.m.LR(0.1)
+            )
+
+        AdaHessian preconditioner can be applied to any other module by passing it to the :code:`inner` argument.
+        Turn off AdaHessian's first momentum to get just the preconditioning. Here is an example of applying
+        AdaHessian preconditioning to nesterov momentum (:code:`tz.m.NAG`):
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.AdaHessian(beta1=0, inner=tz.m.NAG(0.9)),
+                tz.m.LR(0.1)
+            )
+
+    """
+    def __init__(
+        self,
+        beta1: float = 0.9,
+        beta2: float = 0.999,
+        averaging: bool = False,
+        block_size: int | None = 9,
+        update_freq: int = 1,
+        eps: float = 1e-8,
+        hvp_method: Literal['autograd', 'forward', 'central'] = 'autograd',
+        fd_h: float = 1e-3,
+        n_samples = 1,
+        seed: int | None = None,
+        inner: Chainable | None = None
+    ):
+        defaults = dict(beta1=beta1, beta2=beta2, update_freq=update_freq, averaging=averaging, block_size=block_size, eps=eps, hvp_method=hvp_method, n_samples=n_samples, fd_h=fd_h, seed=seed)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, var):
+        params = var.params
+        settings = self.settings[params[0]]
+        hvp_method = settings['hvp_method']
+        fd_h = settings['fd_h']
+        update_freq = settings['update_freq']
+        n_samples = settings['n_samples']
+
+        seed = settings['seed']
+        generator = None
+        if seed is not None:
+            if 'generator' not in self.global_state:
+                self.global_state['generator'] = torch.Generator(params[0].device).manual_seed(seed)
+            generator = self.global_state['generator']
+
+        beta1, beta2, eps, averaging, block_size = self.get_settings(params,
+            'beta1', 'beta2', 'eps', 'averaging', 'block_size', cls=NumberList)
+
+        exp_avg, D_exp_avg_sq = self.get_state(params, 'exp_avg', 'h_exp_avg', cls=TensorList)
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        closure = var.closure
+        assert closure is not None
+
+        D = None
+        if step % update_freq == 0:
+
+            rgrad=None
+            for i in range(n_samples):
+                u = [_rademacher_like(p, generator=generator) for p in params]
+
+                Hvp, rgrad = self.Hvp(u, at_x0=True, var=var, rgrad=rgrad, hvp_method=hvp_method,
+                                     h=fd_h, normalize=True, retain_grad=i < n_samples-1)
+
+                if D is None: D = Hvp
+                else: torch._foreach_add_(D, Hvp)
+
+            assert D is not None
+            if n_samples > 1: torch._foreach_div_(D, n_samples)
+
+            D = TensorList(D).zipmap_args(_block_average, block_size, averaging)
+
+        update = var.get_update()
+        if 'inner' in self.children:
+            update = apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var)
+
+        var.update = adahessian(
+            tensors=TensorList(update),
+            D=TensorList(D) if D is not None else None,
+            exp_avg_=exp_avg,
+            D_exp_avg_sq_=D_exp_avg_sq,
+            beta1=beta1,
+            beta2=beta2,
+            update_freq=update_freq,
+            eps=eps,
+            step=step,
+        )
+        return var

torchzero/modules/optimizers/adam.py

@@ -3,14 +3,14 @@ from functools import partial
 
 import torch
 
-from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList
+from ...core import Module, Target, Transform, apply_transform, Chainable
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 from ..functional import (
     debias, debiased_step_size,
     ema_,
     sqrt_ema_sq_,
 )
-from ..lr.lr import lazy_lr
+from ..step_size.lr import lazy_lr
 from ..momentum.experimental import sqrt_nag_ema_sq_
 from ..momentum.momentum import nag_
 
@@ -27,26 +27,28 @@ def adam_(
     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."""
-    exp_avg_ = ema_(tensors, exp_avg_=exp_avg_, beta=beta1, dampening=0,lerp=True)
 
+    # inner args
+    inner: Module | None = None,
+    params: list[torch.Tensor] | None = None,
+    grads: list[torch.Tensor] | None = None,
+):
+    """Returns new tensors."""
     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)
 
-    if debiased: alpha = debiased_step_size(step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
+    if inner is not None:
+        assert params is not None
+        tensors = TensorList(apply_transform(inner, tensors, params=params, grads=grads))
 
-    # params is None, return update
-    if params_ is None: return (exp_avg_ / sqrt_exp_avg_sq.add_(eps)).lazy_mul(alpha)
+    exp_avg_ = ema_(tensors, exp_avg_=exp_avg_, beta=beta1, dampening=0,lerp=True)
+    if debiased: alpha = debiased_step_size(step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
+    return (exp_avg_.lazy_mul(alpha) / sqrt_exp_avg_sq.add_(eps))
 
-    # update params in-place
-    params_.addcdiv_(exp_avg_, sqrt_exp_avg_sq.add_(eps), -alpha)
-    return None
+class Adam(Transform):
+    """Adam. Divides gradient EMA by EMA of gradient squares with debiased step size.
 
-class Adam(Module):
-    """Adam. Divides gradient EMA by EMA of gradient squares with debiased step size. This implementation is slightly different from
-    pytorch in that debiasing is applied after adding epsilon.
+    This implementation is identical to :code:`torch.optim.Adam`.
 
     Args:
         beta1 (float, optional): momentum. Defaults to 0.9.
@@ -66,36 +68,29 @@ class Adam(Module):
         alpha: float = 1.,
         pow: float = 2,
         debiased: bool = True,
+        inner: Chainable | None = None
     ):
         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')
+        super().__init__(defaults, uses_grad=False)
+
+        if inner is not None: self.set_child('inner', inner)
 
     @torch.no_grad
-    def step(self, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        beta1,beta2,eps,alpha=self.get_settings('beta1','beta2','eps','alpha', params=vars.params, cls=NumberList)
-        amsgrad,pow,debiased = self.getter(self.settings[vars.params[0]])
+        beta1,beta2,eps,alpha=unpack_dicts(settings, 'beta1','beta2','eps','alpha', cls=NumberList)
+        amsgrad,pow,debiased = itemgetter('amsgrad','pow','debiased')(settings[0])
 
         if amsgrad:
-            exp_avg, exp_avg_sq, max_exp_avg_sq = self.get_state('exp_avg','exp_avg_sq','max_exp_avg_sq', params=vars.params, cls=TensorList)
+            exp_avg, exp_avg_sq, max_exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', cls=TensorList)
         else:
-            exp_avg, exp_avg_sq = self.get_state('exp_avg','exp_avg_sq', params=vars.params, cls=TensorList)
+            exp_avg, exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', cls=TensorList)
             max_exp_avg_sq = None
 
-        # if this is last module, update parameters in-place with slightly more efficient addcdiv_
-        if vars.is_last:
-            if vars.last_module_lrs is not None: alpha = alpha * vars.last_module_lrs
-            passed_params = TensorList(vars.params)
-            vars.stop = True
-            vars.skip_update = True
-
-        else:
-            passed_params = None
 
-        vars.update = adam_(
-            tensors=TensorList(vars.get_update()),
+        return adam_(
+            tensors=TensorList(tensors),
             exp_avg_=exp_avg,
             exp_avg_sq_=exp_avg_sq,
             alpha=alpha,
@@ -106,7 +101,10 @@ class Adam(Module):
             pow=pow,
             debiased=debiased,
             max_exp_avg_sq_=max_exp_avg_sq,
-            params_=passed_params,
-        )
 
-        return vars
+            # inner args
+            inner=self.children.get("inner", None),
+            params=params,
+            grads=grads,
+
+        )

torchzero/modules/optimizers/adan.py

@@ -0,0 +1,110 @@
+import torch
+
+from ...core import Transform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+
+def adan_(
+    g: TensorList,
+    g_prev_: TensorList,
+    m_: TensorList, # exponential moving average
+    v_: TensorList, # exponential moving average of gradient differences
+    n_: TensorList, # kinda like squared momentum
+    n_prev_: TensorList | None,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    beta3: float | NumberList,
+    eps: float | NumberList,
+    use_n_prev: bool,
+):
+    """Returns new tensors."""
+    m_.lerp_(g, 1-beta1)
+
+    y = g - g_prev_
+    v_.lerp_(y, 1-beta2)
+
+    y.mul_(1-beta2).add_(g)
+    n_.mul_(beta3).addcmul_(y, y, 1-beta3)
+
+    if use_n_prev:
+        assert n_prev_ is not None
+        ns = n_prev_.clone()
+        n_prev_.copy_(n_)
+        n_ = ns
+
+    eta = n_.sqrt().add_(eps).reciprocal_()
+    term = m_ + (1-beta2)*v_
+    update = eta.mul_(term)
+
+    g_prev_.copy_(g)
+
+    return update
+
+
+class Adan(Transform):
+    """Adaptive Nesterov Momentum Algorithm from https://arxiv.org/abs/2208.06677
+
+    Args:
+        beta1 (float, optional): momentum. Defaults to 0.98.
+        beta2 (float, optional): momentum for gradient differences. Defaults to 0.92.
+        beta3 (float, optional): thrid (squared) momentum. Defaults to 0.99.
+        eps (float, optional): epsilon. Defaults to 1e-8.
+        use_n_prev (bool, optional):
+            whether to use previous gradient differences momentum.
+
+    Reference:
+        Xie, X., Zhou, P., Li, H., Lin, Z., & Yan, S. (2024). Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. IEEE Transactions on Pattern Analysis and Machine Intelligence. https://arxiv.org/abs/2208.06677
+    """
+    def __init__(
+        self,
+        beta1: float = 0.98,
+        beta2: float = 0.92,
+        beta3: float = 0.99,
+        eps: float = 1e-8,
+        use_n_prev: bool = False,
+    ):
+        defaults=dict(beta1=beta1,beta2=beta2,beta3=beta3,eps=eps,use_n_prev=use_n_prev)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        beta1,beta2,beta3,eps=unpack_dicts(settings, 'beta1','beta2','beta3','eps', cls=NumberList)
+        s = settings[0]
+        use_n_prev = s['use_n_prev']
+
+        g_prev, m, v, n = unpack_states(states, tensors, 'g_prev','m','v','n', cls=TensorList)
+
+
+        if use_n_prev:
+            n_prev = unpack_states(states, tensors, 'n_prev', cls=TensorList)
+        else:
+            n_prev = None
+
+        if step == 1:
+            # initial values, also runs on restarts
+            m.copy_(tensors)
+            n.set_(tensors ** 2)
+            v.zero_()
+            g_prev.copy_(tensors)
+            if n_prev is not None: n_prev.set_(tensors ** 2)
+
+        if step == 2:
+            v.set_(tensors - g_prev)
+
+        update = adan_(
+            g=tensors,
+            g_prev_=g_prev,
+            m_=m,
+            v_=v,
+            n_=n,
+            n_prev_=n_prev,
+            beta1=beta1,
+            beta2=beta2,
+            beta3=beta3,
+            eps=eps,
+            use_n_prev=use_n_prev,
+        )
+
+        return update

torchzero/modules/optimizers/adaptive_heavyball.py

@@ -0,0 +1,57 @@
+import torch
+from ...core import  Transform
+from ...utils import TensorList, unpack_dicts, unpack_states
+
+
+def adaptive_heavy_ball(f, f_star, f_prev, g: TensorList, g_prev: TensorList, p: TensorList, p_prev: TensorList):
+    if f - f_star <= torch.finfo(p[0].dtype).eps: return g
+
+    g_g = g.dot(g)
+    g_gp = g.dot(g_prev)
+    num = -(f - f_star) * g.dot(g_prev)
+    denom = (f_prev - f_star) * g_g + (f - f_star) * g_gp
+    m = num/denom
+
+    h = 2*(f - f_star) / g_g
+    return (1 + m) * h * g - m*(p-p_prev)
+
+
+class AdaptiveHeavyBall(Transform):
+    """Adaptive heavy ball from https://hal.science/hal-04832983v1/file/OJMO_2024__5__A7_0.pdf.
+
+    This is related to conjugate gradient methods, it may be very good for non-stochastic convex objectives, but won't work on stochastic ones.
+
+    .. note::
+        The step size is determined by the algorithm, so learning rate modules shouldn't be used.
+
+    Args:
+        f_star (int, optional):
+            (estimated) minimal possible value of the objective function (lowest possible loss). Defaults to 0.
+        tol (float, optional):
+            tolerance on objective value change.
+    """
+    def __init__(self, f_star: float = 0):
+        defaults = dict(f_star=f_star)
+        super().__init__(defaults, uses_grad=False, uses_loss=True)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        assert loss is not None
+        tensors = TensorList(tensors)
+        setting = settings[0]
+        f_star = setting['f_star']
+
+        f_prev = self.global_state.get('f_prev', None)
+        p_prev, g_prev = unpack_states(states, tensors, 'p_prev', 'g_prev', init=[params,tensors], cls=TensorList)
+
+        if f_prev is None:
+            self.global_state['f_prev'] = loss
+            h = 2*(loss - f_star) / tensors.dot(tensors)
+            return h * tensors
+
+        update = adaptive_heavy_ball(f=loss, f_star=f_star, f_prev=f_prev, g=tensors, g_prev=g_prev, p=TensorList(params), p_prev=p_prev)
+
+        self.global_state['f_prev'] = loss
+        p_prev.copy_(params)
+        g_prev.copy_(tensors)
+        return update

torchzero/modules/optimizers/esgd.py

@@ -0,0 +1,171 @@
+import math
+from collections.abc import Callable
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, Target, Transform, apply_transform
+from ...utils import NumberList, TensorList, as_tensorlist
+
+
+def esgd_(
+    tensors_: TensorList,
+    D: TensorList | None,
+    D_sq_acc_: TensorList,
+    damping: float | NumberList,
+    update_freq: int,
+    step: int,
+    i: int,
+):
+    # update preconditioner
+    if step % update_freq == 0:
+        assert D is not None
+        D_sq_acc_.addcmul_(D, D)
+        i += 1
+    else:
+        assert D is None
+
+    denom = (D_sq_acc_ / max(i, 1)).sqrt_().add_(damping)
+    return tensors_.div_(denom), i
+
+
+class ESGD(Module):
+    """Equilibrated Gradient Descent (https://arxiv.org/abs/1502.04390)
+
+    This is similar to Adagrad, but the accumulates squared randomized hessian diagonal estimates instead of squared gradients.
+
+    .. note::
+        In most cases Adagrad should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Adagrad preconditioning to another module's output.
+
+    .. note::
+        If you are using gradient estimators or reformulations, set :code:`hvp_method` to "forward" or "central".
+
+    .. note::
+        This module requires a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients for calculating HVPs.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    Args:
+        damping (float, optional): added to denominator for stability. Defaults to 1e-4.
+        update_freq (int, optional):
+            frequency of updating hessian diagonal estimate via a hessian-vector product.
+            This value can be increased to reduce computational cost. Defaults to 20.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        h (float, optional): finite difference step size if :code:`hvp_method` is "forward" or "central". Defaults to 1e-3.
+        n_samples (int, optional):
+            number of hessian-vector products with random vectors to evaluate each time when updating
+            the preconditioner. Larger values may lead to better hessian diagonal estimate. Defaults to 1.
+        seed (int | None, optional): seed for random vectors. Defaults to None.
+        inner (Chainable | None, optional):
+            Inner module. If this is specified, operations are performed in the following order.
+            1. compute hessian diagonal estimate.
+            2. pass inputs to :code:`inner`.
+            3. momentum and preconditioning are applied to the ouputs of :code:`inner`.
+
+    Examples:
+        Using ESGD:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ESGD(),
+                tz.m.LR(0.1)
+            )
+
+        ESGD preconditioner can be applied to any other module by passing it to the :code:`inner` argument. Here is an example of applying
+        ESGD preconditioning to nesterov momentum (:code:`tz.m.NAG`):
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ESGD(beta1=0, inner=tz.m.NAG(0.9)),
+                tz.m.LR(0.1)
+            )
+
+    """
+    def __init__(
+        self,
+        damping: float = 1e-4,
+        update_freq: int = 20,
+        hvp_method: Literal['autograd', 'forward', 'central'] = 'autograd',
+        fd_h: float = 1e-3,
+        n_samples = 1,
+        seed: int | None = None,
+        inner: Chainable | None = None
+    ):
+        defaults = dict(damping=damping, update_freq=update_freq, hvp_method=hvp_method, n_samples=n_samples, fd_h=fd_h, seed=seed)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, var):
+        params = var.params
+        settings = self.settings[params[0]]
+        hvp_method = settings['hvp_method']
+        fd_h = settings['fd_h']
+        update_freq = settings['update_freq']
+        n_samples = settings['n_samples']
+
+        seed = settings['seed']
+        generator = None
+        if seed is not None:
+            if 'generator' not in self.global_state:
+                self.global_state['generator'] = torch.Generator(params[0].device).manual_seed(seed)
+            generator = self.global_state['generator']
+
+        damping = self.get_settings(params, 'damping', cls=NumberList)
+        D_sq_acc = self.get_state(params, 'D_sq_acc', cls=TensorList)
+        i = self.global_state.get('i', 0)
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        closure = var.closure
+        assert closure is not None
+
+        D = None
+        if step % update_freq == 0:
+
+            rgrad=None
+            for j in range(n_samples):
+                u = [torch.randn(p.size(), generator=generator, device=p.device, dtype=p.dtype) for p in params]
+
+                Hvp, rgrad = self.Hvp(u, at_x0=True, var=var, rgrad=rgrad, hvp_method=hvp_method,
+                                     h=fd_h, normalize=True, retain_grad=j < n_samples-1)
+
+                if D is None: D = Hvp
+                else: torch._foreach_add_(D, Hvp)
+
+            assert D is not None
+            if n_samples > 1: torch._foreach_div_(D, n_samples)
+
+            D = TensorList(D)
+
+        update = var.get_update()
+        if 'inner' in self.children:
+            update = apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var)
+
+        var.update, self.global_state['i'] = esgd_(
+            tensors_=TensorList(update),
+            D=TensorList(D) if D is not None else None,
+            D_sq_acc_=D_sq_acc,
+            damping=damping,
+            update_freq=update_freq,
+            step=step,
+            i=i,
+        )
+        return var