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

torchzero/modules/experimental/absoap.py

@@ -1,12 +1,14 @@
 from operator import itemgetter
+from typing import Literal
 
 import torch
-from typing import Literal
-from ...core import Chainable, Transform, apply
+
+from ...core import Chainable, Transform
 from ..optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
+from ..optimizers.soap import project, project_back, get_orthogonal_matrix, get_orthogonal_matrix_QR
 
 @torch.no_grad
-def update_soap_covariances_(
+def update_absoap_covariances_(
     g1: torch.Tensor,
     g2: torch.Tensor,
     GGs_: list[torch.Tensor | None],
@@ -19,138 +21,36 @@ def update_soap_covariances_(
         if beta is None: GG.add_(torch.tensordot(g1, g2, (axes, axes))) # pyright:ignore[reportArgumentType]
         else: GG.lerp_(torch.tensordot(g1, g2, (axes, axes)), 1-beta) # pyright:ignore[reportArgumentType]
 
-@torch.no_grad
-def project(tensors: torch.Tensor, Q: list[torch.Tensor | None]):
-    """
-    Projects the gradient to the eigenbases of the preconditioner.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat, dims=[[0], [0]]) # pyright:ignore[reportArgumentType]
-        else:
-            # I don't understand this part but it is in https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
 
-@torch.no_grad
-def project_back(tensors: torch.Tensor, Q: list[torch.Tensor| None]):
-    """
-    Projects the gradient back to the original space.
-    """
-    for mat in Q:
-        if mat is None: continue
-        if len(mat) > 0:
-            tensors = torch.tensordot(tensors, mat,dims=[[0], [1]]) # pyright:ignore[reportArgumentType]
-        else:
-            permute_order = list(range(1, len(tensors.shape))) + [0]
-            tensors = tensors.permute(permute_order)
-
-    return tensors
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py
-@torch.no_grad
-def get_orthogonal_matrix(mat: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using torch.linalg.eigh decomposition.
-    """
-    matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m in mat:
-        if m is None: continue
-        if len(m) == 0:
-            matrix.append([])
-            continue
-        if m.dtype != torch.float:
-            original_type = m.dtype
-            original_device = m.device
-            matrix.append(m.float())
-        else:
-            float_data = True
-            matrix.append(m)
-
-    final = []
-    for m in matrix:
-        if len(m) == 0:
-            final.append([])
-            continue
-        try:
-            _, Q = torch.linalg.eigh(m+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-        except Exception:
-            _, Q = torch.linalg.eigh(m.to(torch.float64)+1e-30*torch.eye(m.shape[0], device=m.device)) # pylint:disable=not-callable
-            Q = Q.to(m.dtype)
-        Q = torch.flip(Q, [1])
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-    return final
-
-# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py#L240
-@torch.no_grad
-def get_orthogonal_matrix_QR(exp_avg_sq: torch.Tensor, GG: list[torch.Tensor | None], Q_list: list[torch.Tensor | None]):
-    """
-    Computes the eigenbases of the preconditioner using one round of power iteration
-    followed by torch.linalg.qr decomposition.
-    """
-    matrix = []
-    orth_matrix = []
-    float_data = False
-    original_type = original_device = None
-    for m,o in zip(GG, Q_list):
-        if m is None: continue
-        assert o is not None
-
-        if len(m) == 0:
-            matrix.append([])
-            orth_matrix.append([])
-            continue
-        if m.data.dtype != torch.float:
-            original_type = m.data.dtype
-            original_device = m.data.device
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-        else:
-            float_data = True
-            matrix.append(m.data.float())
-            orth_matrix.append(o.data.float())
-
-    final = []
-    for ind, (m,o) in enumerate(zip(matrix, orth_matrix)):
-        if len(m)==0:
-            final.append([])
-            continue
-        est_eig = torch.diag(o.T @ m @ o)
-        sort_idx = torch.argsort(est_eig, descending=True)
-        exp_avg_sq = exp_avg_sq.index_select(ind, sort_idx)
-        o = o[:,sort_idx]
-        power_iter = m @ o
-        Q, _ = torch.linalg.qr(power_iter) # pylint:disable=not-callable
-
-        if not float_data:
-            Q = Q.to(original_device).type(original_type)
-        final.append(Q)
-
-    return final, exp_avg_sq
-
-Source=Literal['p','g','s','y', 'gy', 'sy', 'sn', 'yn', 'gys', 'sys','sn', 'yn']
+Source=Literal['p','g','s','y', 'gy', 'sy', 'sn', 'yn', 'gys', 'sys']
 class ABSOAP(Transform):
-    """SOAP but with two extra letters included in its name in order to improve converence
-
-    so what you can do is choose what goes into what ,and that is supposed to be good.
+    """SOAP but with some extra options for testing.
+
+    .. warning::
+        This module is just for testing my stupid ideas.
+
+    Args:
+        scale_by_s - whether to scale y by s
+        gg1 - 1st vector into GGᵀ
+        gg2 - 2nd vector into GGᵀ
+        ema1 - vector into 1st momentum
+        ema2 - 2 vectors into 2nd momentum
+        rel1 - if True, multiplies gg1 by params
+        rel2 - same but for gg2
+        norm - if True, gg1 a and gg2 are normalized, and I need to make that into a letter
+
+    letters:
+        p - params
+        g - grad
+        s - param difference
+        y - grad difference
+        gy - g+y
+        sy - s+y
+        sn - s normalized
+        yn - y normalized
+        gys - g + y#g
+        sys - s + y#s
 
-    new args
-
-    scale by s whether to scale gradient differences by parameter differences
-
-    y_to_ema2 whether to use gradient differences for exponential moving average too
-
-    okay I changed these args into another ones
-
-    BASICALLY THIS IS FOR MY EXPERIMENTS
     """
     def __init__(
         self,
@@ -166,8 +66,8 @@ class ABSOAP(Transform):
         alpha: float = 1,
         bias_correction: bool = True,
         scale_by_s: bool = True,
-        first: Source='g',
-        second: Source='g',
+        gg1: Source='g',
+        gg2: Source='g',
         ema1: Source='g',
         ema2: tuple[Source, Source] = ('g','g'),
         rel1: bool=False,
@@ -189,29 +89,27 @@ class ABSOAP(Transform):
             scale_by_s=scale_by_s,
             ema1=ema1,
             ema2=ema2,
-            first=first,
-            second=second,
+            first=gg1,
+            second=gg2,
             rel1=rel1, rel2=rel2,
             norm=norm,
         )
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         updates = []
         # update preconditioners
-        for i,(p,t) in enumerate(zip(params, tensors)):
-            state = self.state[p]
-            settings = self.settings[p]
+        for i,(p,t, state, setting) in enumerate(zip(params, tensors, states, settings)):
             beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps, alpha = itemgetter(
-                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'alpha')(settings)
-            scale_by_s = settings['scale_by_s']
-            ema1 = settings['ema1']
-            ema2 = settings['ema2']
-            first=settings['first']
-            second=settings['second']
-            rel1 = settings['rel1']; rel2 = settings['rel2']
-            norm=settings['norm']
+                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'alpha')(setting)
+            scale_by_s = setting['scale_by_s']
+            ema1 = setting['ema1']
+            ema2 = setting['ema2']
+            first=setting['first']
+            second=setting['second']
+            rel1 = setting['rel1']; rel2 = setting['rel2']
+            norm=setting['norm']
 
             if merge_small:
                 t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
@@ -219,8 +117,8 @@ class ABSOAP(Transform):
             if 'g_prev' not in state:
                 state['p_prev'] = p.clone()
                 state['g_prev'] = t.clone()
-                updates.append(tensors[i].clip(-0.1,0.1))
-                continue
+                # updates.append(tensors[i].clip(-0.1,0.1))
+                # continue
 
             p_prev = state['p_prev']
             g_prev = state['g_prev']
@@ -270,11 +168,10 @@ class ABSOAP(Transform):
                 t1 = t1/torch.linalg.vector_norm(t1).clip(min=1e-8) # pylint:disable=not-callable
                 t2 = t2/torch.linalg.vector_norm(t2).clip(min=1e-8) # pylint:disable=not-callable
 
-
             # initialize state on 1st step
             if 'GG' not in state:
                 state["exp_avg"] = torch.zeros_like(t)
-                state["exp_avg_sq"] = torch.ones_like(t)
+                state["exp_avg_sq"] = torch.zeros_like(t)
 
                 if not precondition_1d and t.ndim <= 1:
                     state['GG'] = []
@@ -287,7 +184,7 @@ class ABSOAP(Transform):
                     state['GG'] = None
 
                 if state['GG'] is not None:
-                    update_soap_covariances_(t1, t2, GGs_=state['GG'], beta=shampoo_beta)
+                    update_absoap_covariances_(t1, t2, GGs_=state['GG'], beta=shampoo_beta)
                     state['Q'] = get_orthogonal_matrix(state['GG'])
 
                 state['step'] = 0
@@ -334,7 +231,7 @@ class ABSOAP(Transform):
             if z1_projected is not None:
                 update = project_back(update, state["Q"])
 
-            if settings['bias_correction']:
+            if setting['bias_correction']:
                 bias_correction1 = 1.0 - beta1 ** (state["step"]+1)
                 bias_correction2 = 1.0 - beta2 ** (state["step"]+1)
                 update *= ((bias_correction2 ** .5) / bias_correction1) * alpha
@@ -349,8 +246,8 @@ class ABSOAP(Transform):
 
             # Update is done after the gradient step to avoid using current gradients in the projection.
             if state['GG'] is not None:
-                update_soap_covariances_(t1, t2, state['GG'], shampoo_beta)
-                if state['step'] % settings['precond_freq'] == 0:
+                update_absoap_covariances_(t1, t2, state['GG'], shampoo_beta)
+                if state['step'] % setting['precond_freq'] == 0:
                     state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, state['GG'], state['Q'])
 
         return updates

torchzero/modules/experimental/adadam.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_
 
@@ -50,7 +50,13 @@ def adadam_(
     return None
 
 class Adadam(Module):
-    """Adam with a diagonally preconditioned preconditioner."""
+    """Adam with a diagonally preconditioned preconditioner.
+
+    Verdict: I haven't tested this yet.
+
+    .. warning::
+        Experimental.
+    """
     def __init__(
         self,
         beta1: float = 0.9,
@@ -67,31 +73,32 @@ class Adadam(Module):
         self.getter = itemgetter('amsgrad','pow','debiased')
 
     @torch.no_grad
-    def step(self, vars):
+    def step(self, var):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+        params = var.params
 
-        beta1,beta2,precond_beta,eps,alpha=self.get_settings('beta1','beta2','precond_beta','eps','alpha', params=vars.params, cls=NumberList)
-        amsgrad,pow,debiased = self.getter(self.settings[vars.params[0]])
+        beta1,beta2,precond_beta,eps,alpha=self.get_settings(params, 'beta1','beta2','precond_beta','eps','alpha', cls=NumberList)
+        amsgrad,pow,debiased = self.getter(self.settings[var.params[0]])
 
         if amsgrad:
-            exp_avg, exp_avg_sq, exp_avg_qu, max_exp_avg_sq, max_exp_avg_qu = self.get_state('exp_avg','exp_avg_sq', 'exp_avg_qu', 'max_exp_avg_sq', 'max_exp_avg_qu', params=vars.params, cls=TensorList)
+            exp_avg, exp_avg_sq, exp_avg_qu, max_exp_avg_sq, max_exp_avg_qu = self.get_state(params, 'exp_avg','exp_avg_sq', 'exp_avg_qu', 'max_exp_avg_sq', 'max_exp_avg_qu', cls=TensorList)
         else:
-            exp_avg, exp_avg_sq, exp_avg_qu = self.get_state('exp_avg','exp_avg_sq', 'exp_avg_qu', params=vars.params, cls=TensorList)
+            exp_avg, exp_avg_sq, exp_avg_qu = self.get_state(params, 'exp_avg','exp_avg_sq', 'exp_avg_qu', cls=TensorList)
             max_exp_avg_sq = None
             max_exp_avg_qu = 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
+        if var.is_last:
+            if var.last_module_lrs is not None: alpha = alpha * var.last_module_lrs
+            passed_params = TensorList(var.params)
+            var.stop = True
+            var.skip_update = True
 
         else:
             passed_params = None
 
-        vars.update = adadam_(
-            tensors=TensorList(vars.get_update()),
+        var.update = adadam_(
+            tensors=TensorList(var.get_update()),
             exp_avg_=exp_avg,
             exp_avg_sq_=exp_avg_sq,
             exp_avg_qu_=exp_avg_qu,
@@ -108,4 +115,4 @@ class Adadam(Module):
             params_=passed_params,
         )
 
-        return vars
+        return var

torchzero/modules/experimental/adamY.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_
 
@@ -64,14 +64,10 @@ def adamy_(
 class AdamY(Module):
     """Adam but uses scaled gradient differences for second momentum.
 
-    Args:
-        beta1 (float, optional): momentum. Defaults to 0.9.
-        beta2 (float, optional): second momentum. Defaults to 0.999.
-        eps (float, optional): epsilon. Defaults to 1e-8.
-        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.
+    Verdict: I haven't tested this yet.
+
+    .. warning::
+        Experimental.
     """
     def __init__(
         self,
@@ -88,36 +84,36 @@ class AdamY(Module):
         self.getter = itemgetter('amsgrad','pow','debiased')
 
     @torch.no_grad
-    def step(self, vars):
+    def step(self, var):
         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=self.get_settings(var.params, 'beta1','beta2','eps','alpha', cls=NumberList)
+        amsgrad,pow,debiased = self.getter(self.settings[var.params[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 = self.get_state(var.params,'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 = self.get_state(var.params, '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
+        if var.is_last:
+            if var.last_module_lrs is not None: alpha = alpha * var.last_module_lrs
+            passed_params = TensorList(var.params)
+            var.stop = True
+            var.skip_update = True
 
         else:
             passed_params = None
 
-        p_prev = self.get_state('p_prev', params=vars.params, cls=TensorList)
-        g_prev = self.get_state('g_prev', params=vars.params, cls=TensorList)
+        p_prev = self.get_state(var.params, 'p_prev', cls=TensorList)
+        g_prev = self.get_state(var.params, 'g_prev', cls=TensorList)
 
 
-        vars.update = adamy_(
-            p=TensorList(vars.params),
+        var.update = adamy_(
+            p=TensorList(var.params),
             p_prev=p_prev,
-            g=TensorList(vars.get_update()),
+            g=TensorList(var.get_update()),
             g_prev=g_prev,
             exp_avg_=exp_avg,
             exp_avg_sq_=exp_avg_sq,
@@ -132,4 +128,4 @@ class AdamY(Module):
             params_=passed_params,
         )
 
-        return vars
+        return var

torchzero/modules/experimental/adam_lambertw.py

@@ -0,0 +1,149 @@
+from operator import itemgetter
+from functools import partial
+import math
+import torch
+
+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 ..step_size.lr import lazy_lr
+from ..momentum.experimental import sqrt_nag_ema_sq_
+from ..momentum.momentum import nag_
+
+
+def _lambertw_newton_raphson(x: TensorList, iterations=5):
+    # z = torch.zeros_like(x)
+    # mask_neg = x < 0
+    # mask_pos = ~mask_neg
+
+    # z[mask_pos] = torch.log(x[mask_pos] + 1.0)
+
+    # x_neg = x[mask_neg]
+    # z_neg = -1.0 + torch.sqrt(2.0 * (1.0 + math.e * x_neg))
+    # z[mask_neg] = z_neg
+
+    # x is always positive
+    z = (x+1).log_()
+    for _ in range(iterations):
+        exp_z = z.exp()
+        numerator = z * exp_z - x
+        denominator = exp_z * (z + 1.0) + 1e-8
+        delta = numerator / denominator
+        z -= delta
+    return z
+
+# https://github.com/gmgeorg/torchlambertw/blob/main/torchlambertw/special.py
+def _lambertw_winitzki(x: TensorList):
+    x_log1p = x.log1p()
+    return x_log1p * (1.0 - x_log1p.log1p() / (2.0 + x_log1p))
+
+
+def adam_lambertw_(
+    tensors: TensorList,
+    exp_avg_: TensorList,
+    exp_avg_xpx_: TensorList,
+    alpha: float | NumberList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    eps: float | NumberList,
+    step: int,
+    pow: float = 2,
+    debiased: bool = True,
+    max_exp_avg_xpx_: TensorList | None = None,
+    iterations: int | None = 5,
+
+    # inner args
+    inner: Module | None = None,
+    params: list[torch.Tensor] | None = None,
+    grads: list[torch.Tensor] | None = None,
+):
+    """Returns new tensors."""
+    tensors_abs = tensors.abs().clip_(max=20)
+    tensors_xpx = tensors_abs.pow_(tensors_abs)
+    exp_avg_xpx_.lerp_(tensors_xpx, 1-beta2)
+
+    if max_exp_avg_xpx_ is not None:
+        max_exp_avg_xpx_.maximum_(exp_avg_xpx_)
+        exp_avg_xpx_ = max_exp_avg_xpx_
+
+    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)
+
+    if iterations is None or iterations < 1: exp_avg_xpx_ = _lambertw_winitzki(exp_avg_xpx_)
+    else: exp_avg_xpx_ = _lambertw_newton_raphson(exp_avg_xpx_, iterations)
+
+    return (exp_avg_.lazy_mul(alpha) / exp_avg_xpx_.add_(eps))
+
+class AdamLambertW(Transform):
+    """Adam but uses abs x^x and LambertW instead of square and sqrt.
+    The gradient will be clipped to 20 because float32 which you have to use otherwise you're PC will explode.
+
+    Args:
+        beta1 (float, optional): momentum. Defaults to 0.9.
+        beta2 (float, optional): second momentum. Defaults to 0.999.
+        eps (float, optional): epsilon. Defaults to 1e-8.
+        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.
+        iterations (int, optional): 0 or None means Winitzki approximation otherwise number of newton raphson iterations.
+    """
+    def __init__(
+        self,
+        beta1: float = 0.9,
+        beta2: float = 0.999,
+        eps: float = 1e-8,
+        amsgrad: bool = False,
+        alpha: float = 1.,
+        pow: float = 2,
+        debiased: bool = True,
+        iterations: int | None = 5,
+        inner: Chainable | None = None
+    ):
+        defaults=dict(beta1=beta1,beta2=beta2,eps=eps,alpha=alpha,amsgrad=amsgrad,pow=pow,debiased=debiased, iterations=iterations)
+        super().__init__(defaults, uses_grad=False)
+
+        if inner is not None: self.set_child('inner', inner)
+
+    @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,iterations = itemgetter('amsgrad','pow','debiased','iterations')(settings[0])
+
+        if amsgrad:
+            exp_avg, exp_avg_xpx, max_exp_avg_xpx = unpack_states(states, tensors, 'exp_avg', 'exp_avg_xpx', 'max_exp_avg_xpx', cls=TensorList)
+        else:
+            exp_avg, exp_avg_xpx = unpack_states(states, tensors, 'exp_avg', 'exp_avg_xpx', cls=TensorList)
+            max_exp_avg_xpx = None
+
+
+        return adam_lambertw_(
+            tensors=TensorList(tensors),
+            exp_avg_=exp_avg,
+            exp_avg_xpx_=exp_avg_xpx,
+            alpha=alpha,
+            beta1=beta1,
+            beta2=beta2,
+            eps=eps,
+            step=step,
+            pow=pow,
+            debiased=debiased,
+            max_exp_avg_xpx_=max_exp_avg_xpx,
+            iterations=iterations,
+
+            # inner args
+            inner=self.children.get("inner", None),
+            params=params,
+            grads=grads,
+
+        )

torchzero/modules/{line_search/trust_region.py → experimental/adaptive_step_size.py}

@@ -2,35 +2,64 @@ from operator import itemgetter
 
 import torch
 
-from .line_search import LineSearch
+from ..line_search import LineSearchBase
 
 
-class TrustRegion(LineSearch):
-    """Basic first order trust region, re-evaluates closure with updated parameters and scales step size based on function value change"""
+class AdaptiveStepSize(LineSearchBase):
+    """Basic first order step size adaptation method. Re-evaluates the function after stepping, if value decreased sufficiently,
+    step size is increased. If value increased, step size is decreased.
+
+    .. note::
+        This works well in some cases, but it is often prone to collapsing.
+        For a more robust alternative use :code:`tz.m.AdaptiveBacktracking`.
+
+    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.
+
+
+    Examples:
+        Adagrad with trust region:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adagrad(),
+                tz.m.TrustRegion()
+            )
+
+    """
     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, vars):
+    def search(self, update, var):
 
-        nplus, nminus, c, init, backtrack, adaptive = itemgetter('nplus','nminus','c','init','backtrack', 'adaptive')(self.settings[vars.params[0]])
+        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, vars, backward=False)
+        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(vars.get_grad(), update))
+        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, vars, backward=False)
+        f_1 = self.evaluate_step_size(step_size, var, backward=False)
 
         proposed = step_size