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

torchzero/modules/ops/utility.py

@@ -6,36 +6,35 @@ from ...core import Module, Target, Transform
 from ...utils.tensorlist import Distributions, TensorList
 
 
-class Clone(Transform):
-    def __init__(self): super().__init__({}, uses_grad=False)
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings): return [t.clone() for t in tensors]
-
-class Grad(Module):
+class Clone(Module):
+    """Clones input. May be useful to store some intermediate result and make sure it doesn't get affected by in-place operations"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
     def step(self, var):
-        var.update = [g.clone() for g in var.get_grad()]
+        var.update = [u.clone() for u in var.get_update()]
         return var
 
-class Params(Module):
+class Grad(Module):
+    """Outputs the gradient"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
     def step(self, var):
-        var.update = [p.clone() for p in var.params]
+        var.update = [g.clone() for g in var.get_grad()]
         return var
 
-class Update(Module):
+class Params(Module):
+    """Outputs parameters"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
     def step(self, var):
-        var.update = [u.clone() for u in var.get_update()]
+        var.update = [p.clone() for p in var.params]
         return var
 
 class Zeros(Module):
+    """Outputs zeros"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
@@ -44,6 +43,7 @@ class Zeros(Module):
         return var
 
 class Ones(Module):
+    """Outputs ones"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
@@ -52,6 +52,7 @@ class Ones(Module):
         return var
 
 class Fill(Module):
+    """Outputs tensors filled with :code:`value`"""
     def __init__(self, value: float):
         defaults = dict(value=value)
         super().__init__(defaults)
@@ -62,6 +63,7 @@ class Fill(Module):
         return var
 
 class RandomSample(Module):
+    """Outputs tensors filled with random numbers from distribution depending on value of :code:`distribution`."""
     def __init__(self, eps: float = 1, distribution: Distributions = 'normal'):
         defaults = dict(eps=eps, distribution=distribution)
         super().__init__(defaults)
@@ -74,6 +76,7 @@ class RandomSample(Module):
         return var
 
 class Randn(Module):
+    """Outputs tensors filled with random numbers from a normal distribution with mean 0 and variance 1."""
     def __init__(self):
         super().__init__({})
 
@@ -83,6 +86,7 @@ class Randn(Module):
         return var
 
 class Uniform(Module):
+    """Outputs tensors filled with random numbers from uniform distribution between :code:`low` and :code:`high`."""
     def __init__(self, low: float, high: float):
         defaults = dict(low=low, high=high)
         super().__init__(defaults)
@@ -94,19 +98,23 @@ class Uniform(Module):
         return var
 
 class GradToNone(Module):
+    """Sets :code:`grad` attribute to None on :code:`var`."""
     def __init__(self): super().__init__()
     def step(self, var):
         var.grad = None
         return var
 
 class UpdateToNone(Module):
+    """Sets :code:`update` attribute to None on :code:`var`."""
     def __init__(self): super().__init__()
     def step(self, var):
         var.update = None
         return var
 
 class Identity(Module):
+    """A placeholder identity operator that is argument-insensitive."""
     def __init__(self, *args, **kwargs): super().__init__()
     def step(self, var): return var
 
-NoOp = Identity
+NoOp = Identity
+"""A placeholder identity operator that is argument-insensitive."""

torchzero/modules/optimizers/__init__.py

@@ -1,7 +1,18 @@
 from .adagrad import Adagrad, FullMatrixAdagrad
+
+# from .curveball import CurveBall
+# from .spectral import SpectralPreconditioner
+from .adahessian import AdaHessian
 from .adam import Adam
+from .adan import Adan
+from .adaptive_heavyball import AdaptiveHeavyBall
+from .esgd import ESGD
+from .ladagrad import LMAdagrad
 from .lion import Lion
+from .mars import MARSCorrection
+from .msam import MSAM, MSAMObjective
 from .muon import DualNormCorrection, MuonAdjustLR, Orthogonalize, orthogonalize_grads_
+from .orthograd import OrthoGrad, orthograd_
 from .rmsprop import RMSprop
 from .rprop import (
     BacktrackOnSignChange,
@@ -10,9 +21,7 @@ from .rprop import (
     SignConsistencyLRs,
     SignConsistencyMask,
 )
+from .sam import ASAM, SAM
 from .shampoo import Shampoo
 from .soap import SOAP
-from .orthograd import OrthoGrad, orthograd_
 from .sophia_h import SophiaH
-# from .curveball import CurveBall
-# from .spectral import SpectralPreconditioner

torchzero/modules/optimizers/adagrad.py

@@ -25,6 +25,7 @@ def adagrad_(
     step: int,
     pow: float = 2,
     use_sqrt: bool = True,
+    divide: bool = False,
 
     # inner args
     inner: Module | None = None,
@@ -40,6 +41,8 @@ def adagrad_(
         assert params is not None
         tensors_ = TensorList(apply_transform(inner, tensors_, params=params, grads=grads))
 
+    if divide: sq_sum_ = sq_sum_ / max(step, 1)
+
     if use_sqrt: tensors_.div_(root(sq_sum_, p=pow, inplace=False).add_(eps)).mul_(clr)
     else: tensors_.div_(sq_sum_.add(eps)).mul_(clr)
 
@@ -48,7 +51,9 @@ def adagrad_(
 
 
 class Adagrad(Transform):
-    """Adagrad, divides by sum of past squares of gradients, matches pytorch Adagrad.
+    """Adagrad, divides by sum of past squares of gradients.
+
+    This implementation is identical to :code:`torch.optim.Adagrad`.
 
     Args:
         lr_decay (float, optional): learning rate decay. Defaults to 0.
@@ -67,23 +72,24 @@ class Adagrad(Transform):
         alpha: float = 1,
         pow: float = 2,
         use_sqrt: bool = True,
+        divide: bool=False,
         inner: Chainable | None = None,
     ):
         defaults = dict(alpha = alpha, lr_decay = lr_decay, initial_accumulator_value=initial_accumulator_value,
-                        eps = eps, pow=pow, use_sqrt = use_sqrt)
+                        eps = eps, pow=pow, use_sqrt = use_sqrt, divide=divide)
         super().__init__(defaults=defaults, uses_grad=False)
 
         if inner is not None:
             self.set_child('inner', inner)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    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
 
         lr_decay,alpha,eps = unpack_dicts(settings, 'lr_decay', 'alpha', 'eps', cls=NumberList)
 
-        pow, use_sqrt = itemgetter('pow', 'use_sqrt')(settings[0])
+        pow, use_sqrt, divide = itemgetter('pow', 'use_sqrt', 'divide')(settings[0])
 
         sq_sum = unpack_states(states, tensors, 'sq_sum', cls=TensorList)
 
@@ -100,6 +106,7 @@ class Adagrad(Transform):
             step=self.global_state["step"],
             pow=pow,
             use_sqrt=use_sqrt,
+            divide=divide,
 
             # inner args
             inner=self.children.get("inner", None),
@@ -110,17 +117,17 @@ class Adagrad(Transform):
 
 
 class FullMatrixAdagrad(TensorwiseTransform):
-    def __init__(self, beta: float | None = None, decay: float | None = None, sqrt:bool=True, concat_params=False, update_freq=1, init: Literal['identity', 'zeros', 'ones', 'GGT'] = 'identity', inner: Chainable | None = None):
-        defaults = dict(beta=beta, decay=decay, sqrt=sqrt, init=init)
-        super().__init__(defaults, uses_grad=False, concat_params=concat_params, update_freq=update_freq, inner=inner)
+    def __init__(self, beta: float | None = None, decay: float | None = None, sqrt:bool=True, concat_params=True, update_freq=1, init: Literal['identity', 'zeros', 'ones', 'GGT'] = 'identity', divide: bool=False, inner: Chainable | None = None):
+        defaults = dict(beta=beta, decay=decay, sqrt=sqrt, init=init, divide=divide)
+        super().__init__(defaults, uses_grad=False, concat_params=concat_params, update_freq=update_freq, inner=inner,)
 
     @torch.no_grad
-    def update_tensor(self, tensor, param, grad, loss, state, settings):
+    def update_tensor(self, tensor, param, grad, loss, state, setting):
         G = tensor.ravel()
         GG = torch.outer(G, G)
-        decay = settings['decay']
-        beta = settings['beta']
-        init = settings['init']
+        decay = setting['decay']
+        beta = setting['beta']
+        init = setting['init']
 
         if 'GG' not in state:
             if init == 'identity': state['GG'] = torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype)
@@ -132,11 +139,14 @@ class FullMatrixAdagrad(TensorwiseTransform):
 
         if beta is not None: state['GG'].lerp_(GG, 1-beta)
         else: state['GG'].add_(GG)
+        state['i'] = state.get('i', 0) + 1 # number of GGTs in sum
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
         GG = state['GG']
-        sqrt = settings['sqrt']
+        sqrt = setting['sqrt']
+        divide = setting['divide']
+        if divide: GG = GG/state.get('i', 1)
 
         if tensor.numel() == 1:
             GG = GG.squeeze()

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

@@ -10,7 +10,7 @@ from ..functional import (
     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_
 
@@ -33,7 +33,7 @@ def adam_(
     params: list[torch.Tensor] | None = None,
     grads: list[torch.Tensor] | None = None,
 ):
-    """Returns new tensors or updates params in-place."""
+    """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)
 
@@ -43,11 +43,12 @@ def adam_(
 
     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_ / sqrt_exp_avg_sq.add_(eps)).lazy_mul(alpha)
+    return (exp_avg_.lazy_mul(alpha) / sqrt_exp_avg_sq.add_(eps))
 
 class Adam(Transform):
-    """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.
+    """Adam. Divides gradient EMA by EMA of gradient squares with debiased step size.
+
+    This implementation is identical to :code:`torch.optim.Adam`.
 
     Args:
         beta1 (float, optional): momentum. Defaults to 0.9.
@@ -75,7 +76,7 @@ class Adam(Transform):
         if inner is not None: self.set_child('inner', inner)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    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=unpack_dicts(settings, 'beta1','beta2','eps','alpha', cls=NumberList)

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