torchzero 0.3.15__py3-none-any.whl → 0.4.1__py3-none-any.whl

torchzero/modules/adaptive/adahessian.py

@@ -3,9 +3,8 @@ from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, Target, Transform, apply_transform
-from ...utils import NumberList, TensorList, as_tensorlist
-from ..functional import debiased_step_size
+from ...core import Chainable, Transform, HVPMethod
+from ...utils import NumberList, TensorList, Distributions, unpack_dicts, unpack_states
 
 def _full_average(hvp: torch.Tensor):
     if hvp.ndim >= 3:  # Conv kernel
@@ -37,41 +36,7 @@ def _block_average(x: torch.Tensor, block_size: int | None, enable: bool):
     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,
-    hessian_power: float | NumberList,
-    step: int,
-):
-    # momentum
-    exp_avg_.lerp_(tensors, 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_.sqrt().pow_(hessian_power).add_(eps)
-    num = exp_avg_ * debiased_step_size(step+1, beta1, beta2)
-
-    return num.div_(denom)
-
-
-class AdaHessian(Module):
+class AdaHessian(Transform):
     """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.
@@ -79,8 +44,6 @@ class AdaHessian(Module):
     Notes:
         - In most cases AdaHessian should be the first module in the chain because it relies on autograd. Use the ``inner`` argument if you wish to apply AdaHessian preconditioning to another module's output.
 
-        - If you are using gradient estimators or reformulations, set ``hvp_method`` to "forward" or "central".
-
         - 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:
@@ -97,17 +60,17 @@ class AdaHessian(Module):
         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".
-        fd_h (float, optional): finite difference step size if ``hvp_method`` is "forward" or "central". Defaults to 1e-3.
+            Determines how hessian-vector products are computed.
+
+            - ``"batched_autograd"`` - uses autograd with batched hessian-vector products. If a single hessian-vector is evaluated, equivalent to ``"autograd"``. Faster than ``"autograd"`` but uses more memory.
+            - ``"autograd"`` - uses autograd hessian-vector products. If multiple hessian-vector products are evaluated, uses a for-loop. Slower than ``"batched_autograd"`` but uses less memory.
+            - ``"fd_forward"`` - uses gradient finite difference approximation with a less accurate forward formula which requires one extra gradient evaluation per hessian-vector product.
+            - ``"fd_central"`` - uses gradient finite difference approximation with a more accurate central formula which requires two gradient evaluations per hessian-vector product.
+
+            Defaults to ``"autograd"``.
+        h (float, optional):
+            The step size for finite difference if ``hvp_method`` is
+            ``"fd_forward"`` or ``"fd_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.
@@ -123,7 +86,7 @@ class AdaHessian(Module):
     Using AdaHessian:
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.AdaHessian(),
         tz.m.LR(0.1)
@@ -134,7 +97,7 @@ class AdaHessian(Module):
     Turn off AdaHessian's first momentum to get just the preconditioning. Here is an example of applying
     AdaHessian preconditioning to nesterov momentum (``tz.m.NAG``):
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.AdaHessian(beta1=0, inner=tz.m.NAG(0.9)),
         tz.m.LR(0.1)
@@ -151,74 +114,82 @@ class AdaHessian(Module):
         update_freq: int = 1,
         eps: float = 1e-8,
         hessian_power: float = 1,
-        hvp_method: Literal['autograd', 'forward', 'central'] = 'autograd',
-        fd_h: float = 1e-3,
+        distribution: Distributions = 'rademacher',
+        hvp_method: HVPMethod = 'autograd',
+        h: float = 1e-3,
         n_samples = 1,
+        zHz: bool = True,
+        debias: bool = True,
         seed: int | None = None,
-        inner: Chainable | None = None
+
+        exp_avg_tfm: Chainable | None = None,
+        D_exp_avg_sq_tfm: Chainable | None = None,
     ):
-        defaults = dict(beta1=beta1, beta2=beta2, update_freq=update_freq, averaging=averaging, block_size=block_size, eps=eps, hessian_power=hessian_power, hvp_method=hvp_method, n_samples=n_samples, fd_h=fd_h, seed=seed)
+        defaults = locals().copy()
+        del defaults['self'], defaults["exp_avg_tfm"], defaults["D_exp_avg_sq_tfm"]
         super().__init__(defaults)
 
-        if inner is not None:
-            self.set_child('inner', inner)
+        self.set_child('exp_avg', exp_avg_tfm)
+        self.set_child('D_exp_avg_sq', D_exp_avg_sq_tfm)
 
     @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']
+    def update_states(self, objective, states, settings):
+        params = objective.params
 
-        seed = settings['seed']
-        generator = self.get_generator(params[0].device, seed)
+        beta1, beta2, averaging, block_size = unpack_dicts(settings, 'beta1', 'beta2', 'averaging', 'block_size', cls=NumberList)
 
-        beta1, beta2, eps, averaging, block_size, hessian_power = self.get_settings(params,
-            'beta1', 'beta2', 'eps', 'averaging', 'block_size', "hessian_power", cls=NumberList)
+        exp_avg, D_exp_avg_sq = unpack_states(states, params, 'exp_avg', 'D_exp_avg_sq', cls=TensorList)
 
-        exp_avg, D_exp_avg_sq = self.get_state(params, 'exp_avg', 'h_exp_avg', cls=TensorList)
+        # ---------------------------- hutchinson hessian ---------------------------- #
+        fs = settings[0]
+        step = self.increment_counter("step", start=0) # 0 on 1st update
+        update_freq = fs['update_freq']
 
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
+        if step % update_freq == 0:
+            self.increment_counter("num_Ds", start=1)
+
+            D, _ = objective.hutchinson_hessian(
+                rgrad = None,
+                at_x0 = True,
+                n_samples = fs['n_samples'],
+                distribution = fs['distribution'],
+                hvp_method = fs['hvp_method'],
+                h = fs['h'],
+                zHz = fs["zHz"],
+                generator = self.get_generator(params[0].device, fs["seed"]),
+            )
 
-        closure = var.closure
-        assert closure is not None
+            D = TensorList(D).zipmap_args(_block_average, block_size, averaging)
+            D_exp_avg_sq.mul_(beta2).addcmul_(D, D, value=1-beta2)
 
-        D = None
-        if step % update_freq == 0:
+        # --------------------------------- momentum --------------------------------- #
+        tensors = objective.get_updates() # do this after hutchinson to not disturb autograd
+        exp_avg.lerp_(tensors, 1-beta1)
 
-            rgrad=None
-            for i in range(n_samples):
-                u = [_rademacher_like(p, generator=generator) for p in params]
 
-                Hvp, rgrad = var.hessian_vector_product(u, at_x0=True, rgrad=rgrad, hvp_method=hvp_method,
-                                     h=fd_h, normalize=True, retain_graph=i < n_samples-1)
-                Hvp = tuple(Hvp)
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        params = objective.params
 
-                if D is None: D = Hvp
-                else: torch._foreach_add_(D, Hvp)
+        beta1, beta2, eps, hessian_power = unpack_dicts(settings, 'beta1', 'beta2', 'eps', 'hessian_power', cls=NumberList)
+        exp_avg, D_exp_avg_sq = unpack_states(states, params, 'exp_avg', 'D_exp_avg_sq', cls=TensorList)
 
-            assert D is not None
-            if n_samples > 1: torch._foreach_div_(D, n_samples)
+        # ---------------------------------- debias ---------------------------------- #
+        if settings[0]["debias"]:
+            bias_correction1 = 1.0 - (beta1 ** (self.global_state["step"] + 1))
+            bias_correction2 = 1.0 - (beta2 ** self.global_state["num_Ds"])
+            exp_avg = exp_avg / bias_correction1
+            D_exp_avg_sq = D_exp_avg_sq / bias_correction2
 
-            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,
-            hessian_power=hessian_power,
-            step=step,
-        )
-        return var
+        # -------------------------------- transforms -------------------------------- #
+        exp_avg = TensorList(self.inner_step_tensors(
+            "exp_avg", tensors=exp_avg, clone=True, objective=objective, must_exist=False))
+
+        D_exp_avg_sq = TensorList(self.inner_step_tensors(
+            "D_exp_avg_sq", tensors=D_exp_avg_sq, clone=True, objective=objective, must_exist=False))
+
+        # ------------------------------ compute update ------------------------------ #
+        denom = D_exp_avg_sq.lazy_pow(hessian_power / 2) + eps
+        objective.updates = exp_avg / denom
+        return objective

torchzero/modules/adaptive/adam.py

@@ -1,48 +1,11 @@
-from operator import itemgetter
-from functools import partial
-
 import torch
 
-from ...core import Module, Target, Transform, apply_transform, Chainable
+from ...core import Chainable, Module, TensorTransform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
-from ..functional import (
-    debias, debiased_step_size,
-    ema_,
-    sqrt_ema_sq_,
-)
-
-
-def adam_(
-    tensors: TensorList,
-    exp_avg_: TensorList,
-    exp_avg_sq_: TensorList,
-    alpha: float | NumberList,
-    beta1: float | NumberList,
-    beta2: float | NumberList,
-    eps: float | NumberList,
-    step: int,
-    pow: float = 2,
-    debiased: bool = True,
-    max_exp_avg_sq_: TensorList | None = None,
-
-    # 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 inner is not None:
-        assert params is not None
-        tensors = TensorList(apply_transform(inner, tensors, params=params, grads=grads))
-
-    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))
-
-class Adam(Transform):
+from ..opt_utils import debiased_step_size
+
+
+class Adam(TensorTransform):
     """Adam. Divides gradient EMA by EMA of gradient squares with debiased step size.
 
     This implementation is identical to :code:`torch.optim.Adam`.
@@ -54,7 +17,7 @@ class Adam(Transform):
         alpha (float, optional): learning rate. Defaults to 1.
         amsgrad (bool, optional): Whether to divide by maximum of EMA of gradient squares instead. Defaults to False.
         pow (float, optional): power used in second momentum power and root. Defaults to 2.
-        debiased (bool, optional): whether to apply debiasing to momentums based on current step. Defaults to True.
+        debias (bool, optional): whether to apply debiasing to momentums based on current step. Defaults to True.
     """
     def __init__(
         self,
@@ -63,45 +26,59 @@ class Adam(Transform):
         eps: float = 1e-8,
         amsgrad: bool = False,
         alpha: float = 1.,
-        pow: float = 2,
-        debiased: bool = True,
-        inner: Chainable | None = None
+        debias: bool = True,
+
+        exp_avg_tfm: Chainable | None = None,
+        exp_avg_sq_tfm: Chainable | None = None,
     ):
-        defaults=dict(beta1=beta1,beta2=beta2,eps=eps,alpha=alpha,amsgrad=amsgrad,pow=pow,debiased=debiased)
-        super().__init__(defaults, uses_grad=False)
+        defaults = locals().copy()
+        del defaults['self'], defaults["exp_avg_tfm"], defaults["exp_avg_sq_tfm"]
+        super().__init__(defaults)
 
-        if inner is not None: self.set_child('inner', inner)
+        self.set_child('exp_avg', exp_avg_tfm)
+        self.set_child('exp_avg_sq', exp_avg_sq_tfm)
 
     @torch.no_grad
-    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)
-        amsgrad,pow,debiased = itemgetter('amsgrad','pow','debiased')(settings[0])
-
-        if amsgrad:
-            exp_avg, exp_avg_sq, max_exp_avg_sq = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', 'max_exp_avg_sq', cls=TensorList)
+    def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
+        self.increment_counter("step", start=0)
+        beta1, beta2 = unpack_dicts(settings, 'beta1','beta2', cls=NumberList)
+
+        # ----------------------------- initialize states ---------------------------- #
+        if settings[0]["amsgrad"]:
+            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 = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', cls=TensorList)
             max_exp_avg_sq = None
 
+        # ------------------------------ update moments ------------------------------ #
+        exp_avg.lerp_(tensors, weight=1-beta1)
+        exp_avg_sq.mul_(beta2).addcmul_(tensors, tensors, value=1-beta2)
+
+        if max_exp_avg_sq is not None:
+            max_exp_avg_sq.maximum_(exp_avg_sq)
+
+    @torch.no_grad
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state["step"] # 0 on 1st step
+        fs = settings[0]
+
+        if fs["amsgrad"]: key = "max_exp_avg_sq"
+        else: key = "exp_avg_sq"
+        exp_avg, exp_avg_sq = unpack_states(states, tensors, 'exp_avg', key, cls=TensorList)
+        beta1, beta2, alpha, eps = unpack_dicts(settings, 'beta1', 'beta2', 'alpha', 'eps', cls=NumberList)
+
+        # -------------------------------- transforms -------------------------------- #
+        exp_avg = TensorList(self.inner_step_tensors(
+            "exp_avg", tensors=exp_avg, clone=True, params=params, grads=grads, loss=loss, must_exist=False))
+
+        exp_avg_sq = TensorList(self.inner_step_tensors(
+            "exp_avg_sq", tensors=exp_avg_sq, clone=True, params=params, grads=grads, loss=loss, must_exist=False))
+
+        # ---------------------------------- debias ---------------------------------- #
+        if fs["debias"]:
+            alpha = debiased_step_size((step + 1), beta1=beta1, beta2=beta2, alpha=alpha)
+            exp_avg = exp_avg * alpha
 
-        return adam_(
-            tensors=TensorList(tensors),
-            exp_avg_=exp_avg,
-            exp_avg_sq_=exp_avg_sq,
-            alpha=alpha,
-            beta1=beta1,
-            beta2=beta2,
-            eps=eps,
-            step=step,
-            pow=pow,
-            debiased=debiased,
-            max_exp_avg_sq_=max_exp_avg_sq,
-
-            # inner args
-            inner=self.children.get("inner", None),
-            params=params,
-            grads=grads,
-
-        )
+        # ---------------------------------- update ---------------------------------- #
+        return exp_avg / exp_avg_sq.sqrt().add_(eps)

torchzero/modules/adaptive/adan.py

@@ -1,9 +1,9 @@
 import torch
 
-from ...core import Transform
+from ...core import TensorTransform, Chainable
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 
-def adan_(
+def adan_update_(
     g: TensorList,
     g_prev_: TensorList,
     m_: TensorList, # exponential moving average
@@ -12,10 +12,8 @@ def adan_(
     beta1: float | NumberList,
     beta2: float | NumberList,
     beta3: float | NumberList,
-    eps: float | NumberList,
     step: int,
 ):
-    """Returns new tensors"""
     m_.lerp_(g, 1 - beta1)
 
     if step == 1:
@@ -26,7 +24,18 @@ def adan_(
         term = g + beta2 * diff
 
     n_.mul_(beta3).addcmul_(term, term, value=(1 - beta3))
+    g_prev_.copy_(g)
 
+def adan_apply_(
+    m_: TensorList, # exponential moving average
+    v_: TensorList, # exponential moving average of gradient differences
+    n_: TensorList, # kinda like squared momentum
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    beta3: float | NumberList,
+    eps: float | NumberList,
+    step: int,
+):
     m = m_ / (1.0 - beta1**step)
     v = v_ / (1.0 - beta2**step)
     n = n_ / (1.0 - beta3**step)
@@ -35,13 +44,12 @@ def adan_(
     num = m + beta2 * v
 
     update = num.div_(denom)
-    g_prev_.copy_(g)
 
     return update
 
 
 
-class Adan(Transform):
+class Adan(TensorTransform):
     """Adaptive Nesterov Momentum Algorithm from https://arxiv.org/abs/2208.06677
 
     Args:
@@ -49,18 +57,17 @@ class Adan(Transform):
         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.
 
     Example:
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adan(),
         tz.m.LR(1e-3),
     )
+    ```
     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
+        [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,
@@ -68,29 +75,41 @@ class Adan(Transform):
         beta2: float = 0.92,
         beta3: float = 0.99,
         eps: float = 1e-8,
+
+        m_tfm: Chainable | None = None,
+        v_tfm: Chainable | None = None,
+        n_tfm: Chainable | None = None,
     ):
-        defaults=dict(beta1=beta1,beta2=beta2,beta3=beta3,eps=eps)
+        defaults=dict(beta1=beta1, beta2=beta2, beta3=beta3, eps=eps)
         super().__init__(defaults, uses_grad=False)
 
+        self.set_child("m", m_tfm)
+        self.set_child("v", v_tfm)
+        self.set_child("n", n_tfm)
+
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_update(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)
-        g_prev, m, v, n = unpack_states(states, tensors, 'g_prev','m','v','n', cls=TensorList)
-
-        update = adan_(
-            g=tensors,
-            g_prev_=g_prev,
-            m_=m,
-            v_=v,
-            n_=n,
-            beta1=beta1,
-            beta2=beta2,
-            beta3=beta3,
-            eps=eps,
-            step=step,
-        )
-
-        return update
+        step = self.increment_counter("step", start=0)
+
+        beta1, beta2, beta3 = unpack_dicts(settings, 'beta1','beta2','beta3', cls=NumberList)
+        g_prev, m, v, n = unpack_states(states, tensors, 'g_prev', 'm', 'v', 'n', cls=TensorList)
+
+        adan_update_(g=tensors, g_prev_=g_prev, m_=m, v_=v, n_=n, beta1=beta1, beta2=beta2, beta3=beta3, step=step+1)
+
+    @torch.no_grad
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        step = self.global_state["step"] # 0 on 1st step
+
+        beta1, beta2, beta3, eps = unpack_dicts(settings, 'beta1','beta2','beta3', 'eps', cls=NumberList)
+        m, v, n = unpack_states(states, tensors, 'm', 'v', 'n')
+
+        # -------------------------------- transforms -------------------------------- #
+        m = TensorList(self.inner_step_tensors("m", m, clone=True, params=params, grads=grads, loss=loss, must_exist=False))
+        v = TensorList(self.inner_step_tensors("v", v, clone=True, params=params, grads=grads, loss=loss, must_exist=False))
+        n = TensorList(self.inner_step_tensors("n", n, clone=True, params=params, grads=grads, loss=loss, must_exist=False))
+
+        # ---------------------------------- update ---------------------------------- #
+        return adan_apply_(m_=m, v_=v, n_=n, beta1=beta1, beta2=beta2, beta3=beta3, eps=eps, step=step+1)
+

torchzero/modules/adaptive/adaptive_heavyball.py

@@ -1,5 +1,5 @@
 import torch
-from ...core import  Transform
+from ...core import  TensorTransform
 from ...utils import TensorList, unpack_dicts, unpack_states
 
 
@@ -16,10 +16,10 @@ def adaptive_heavy_ball(f, f_star, f_prev, g: TensorList, g_prev: TensorList, p:
     return (1 + m) * h * g - m*(p-p_prev)
 
 
-class AdaptiveHeavyBall(Transform):
+class AdaptiveHeavyBall(TensorTransform):
     """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.
+    Suitable for quadratic objectives with known f* (loss at minimum).
 
     note:
         The step size is determined by the algorithm, so learning rate modules shouldn't be used.
@@ -30,25 +30,30 @@ class AdaptiveHeavyBall(Transform):
     """
     def __init__(self, f_star: float = 0):
         defaults = dict(f_star=f_star)
-        super().__init__(defaults, uses_grad=False, uses_loss=True)
+        super().__init__(defaults, uses_loss=True)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         assert loss is not None
         tensors = TensorList(tensors)
-        f_star = self.defaults['f_star']
+        f_star = settings[0]['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)
 
+        # -------------------------------- first step -------------------------------- #
         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)
+        # ------------------------------- further steps ------------------------------ #
+        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)
 
+        # --------------------------- store previous values -------------------------- #
         self.global_state['f_prev'] = loss
         p_prev.copy_(params)
         g_prev.copy_(tensors)
+
         return update

torchzero/modules/adaptive/aegd.py

@@ -2,17 +2,18 @@ import math
 
 import torch
 
-from ...core import Transform
+from ...core import TensorTransform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 
 # i've verified, it is identical to official
 # https://github.com/txping/AEGD/blob/master/aegd.py
+# TODO: add a test
 def aegd_(f: torch.Tensor | float, g: TensorList, r_: TensorList, c:float|NumberList=1, eta:float|NumberList=0.1) -> TensorList:
     v = g / (2 * (f + c)**0.5)
     r_ /= 1 + (v ** 2).mul_(2*eta) # update energy
     return 2*eta * r_*v # pyright:ignore[reportReturnType]
 
-class AEGD(Transform):
+class AEGD(TensorTransform):
     """AEGD (Adaptive gradient descent with energy) from https://arxiv.org/abs/2010.05109#page=10.26.
 
     Note:
@@ -20,28 +21,27 @@ class AEGD(Transform):
         To avoid compounding learning rate mofications, remove the ``tz.m.LR`` module if you had it.
 
     Args:
-        eta (float, optional): step size. Defaults to 0.1.
-        c (float, optional): c. Defaults to 1.
-        beta3 (float, optional): thrid (squared) momentum. Defaults to 0.1.
-        eps (float, optional): epsilon. Defaults to 1e-8.
-        use_n_prev (bool, optional):
-            whether to use previous gradient differences momentum.
+        lr (float, optional): learning rate (default: 0.1)
+        c (float, optional): term added to the original objective function (default: 1)
+
+    Reference:
+        [Liu, Hailiang, and Xuping Tian. "AEGD: Adaptive gradient descent with energy." arXiv preprint arXiv:2010.05109 (2020).](https://arxiv.org/pdf/2010.05109)
     """
     def __init__(
         self,
         lr: float = 0.1,
         c: float = 1,
     ):
-        defaults=dict(c=c,lr=lr)
+        defaults = dict(c=c, lr=lr)
         super().__init__(defaults, uses_loss=True)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         assert loss is not None
         tensors = TensorList(tensors)
 
-        c,lr=unpack_dicts(settings, 'c','lr', cls=NumberList)
-        r = unpack_states(states, tensors, 'r', init=lambda t: torch.full_like(t, float(loss+c[0])**0.5), cls=TensorList)
+        c, lr = unpack_dicts(settings, 'c', 'lr', cls=NumberList)
+        r = unpack_states(states, tensors, 'r', init=lambda t: torch.full_like(t, float(loss + c[0])**0.5), cls=TensorList)
 
         update = aegd_(
             f=loss,