torchzero 0.3.14__py3-none-any.whl → 0.4.0__py3-none-any.whl

torchzero/modules/smoothing/laplacian.py

@@ -4,7 +4,7 @@ from collections.abc import Iterable
 import torch
 
 from ...utils.tensorlist import TensorList
-from ...core import Transform, Target
+from ...core import TensorTransform
 
 
 def vector_laplacian_smoothing(input: torch.Tensor, sigma: float = 1) -> torch.Tensor:
@@ -55,7 +55,7 @@ def _precompute_denominator(tensor: torch.Tensor, sigma) -> torch.Tensor:
     v[-1] = 1
     return 1 - sigma * torch.fft.fft(v) # pylint: disable = not-callable
 
-class LaplacianSmoothing(Transform):
+class LaplacianSmoothing(TensorTransform):
     """Applies laplacian smoothing via a fast Fourier transform solver which can improve generalization.
 
     Args:
@@ -70,29 +70,30 @@ class LaplacianSmoothing(Transform):
             what to set on var.
 
     Examples:
-        Laplacian Smoothing Gradient Descent optimizer as in the paper
+    Laplacian Smoothing Gradient Descent optimizer as in the paper
 
-        .. code-block:: python
+    ```python
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.LaplacianSmoothing(),
-                tz.m.LR(1e-2),
-            )
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.LaplacianSmoothing(),
+        tz.m.LR(1e-2),
+    )
+    ```
 
     Reference:
         Osher, S., Wang, B., Yin, P., Luo, X., Barekat, F., Pham, M., & Lin, A. (2022). Laplacian smoothing gradient descent. Research in the Mathematical Sciences, 9(3), 55.
 
     """
-    def __init__(self, sigma:float = 1, layerwise=True, min_numel = 4, target: Target = 'update'):
+    def __init__(self, sigma:float = 1, layerwise=True, min_numel = 4):
         defaults = dict(sigma = sigma, layerwise=layerwise, min_numel=min_numel)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
         # precomputed denominator for when layerwise=False
         self.global_state['full_denominator'] = None
 
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         layerwise = settings[0]['layerwise']
 
         # layerwise laplacian smoothing

torchzero/modules/smoothing/sampling.py

@@ -7,14 +7,15 @@ from typing import Literal, cast
 
 import torch
 
-from ...core import Chainable, Modular, Module, Var
+from ...core import Chainable, Modular, Module, Objective
 from ...core.reformulation import Reformulation
 from ...utils import Distributions, NumberList, TensorList
 from ..termination import TerminationCriteriaBase, make_termination_criteria
 
 
-def _reset_except_self(optimizer: Modular, var: Var, self: Module):
-    for m in optimizer.unrolled_modules:
+def _reset_except_self(objective: Objective, modules, self: Module):
+    assert objective.modular is not None
+    for m in objective.modular.flat_modules:
         if m is not self:
             m.reset()
 
@@ -98,15 +99,15 @@ class GradientSampling(Reformulation):
             self.set_child('termination', make_termination_criteria(extra=termination))
 
     @torch.no_grad
-    def pre_step(self, var):
-        params = TensorList(var.params)
+    def pre_step(self, objective):
+        params = TensorList(objective.params)
 
         fixed = self.defaults['fixed']
 
         # check termination criteria
         if 'termination' in self.children:
             termination = cast(TerminationCriteriaBase, self.children['termination'])
-            if termination.should_terminate(var):
+            if termination.should_terminate(objective):
 
                 # decay sigmas
                 states = [self.state[p] for p in params]
@@ -118,7 +119,7 @@ class GradientSampling(Reformulation):
 
                 # reset on sigmas decay
                 if self.defaults['reset_on_termination']:
-                    var.post_step_hooks.append(partial(_reset_except_self, self=self))
+                    objective.post_step_hooks.append(partial(_reset_except_self, self=self))
 
                 # clear perturbations
                 self.global_state.pop('perts', None)
@@ -136,7 +137,7 @@ class GradientSampling(Reformulation):
             self.global_state['perts'] = perts
 
     @torch.no_grad
-    def closure(self, backward, closure, params, var):
+    def closure(self, backward, closure, params, objective):
         params = TensorList(params)
         loss_agg = None
         grad_agg = None
@@ -160,7 +161,7 @@ class GradientSampling(Reformulation):
 
         # evaluate at x_0
         if include_x0:
-            f_0 = cast(torch.Tensor, var.get_loss(backward=backward))
+            f_0 = objective.get_loss(backward=backward)
 
             isfinite = math.isfinite(f_0)
             if isfinite:
@@ -168,7 +169,7 @@ class GradientSampling(Reformulation):
                 loss_agg = f_0
 
             if backward:
-                g_0 = var.get_grad()
+                g_0 = objective.get_grads()
                 if isfinite: grad_agg = g_0
 
         # evaluate at x_0 + p for each perturbation

torchzero/modules/step_size/adaptive.py

@@ -5,9 +5,9 @@ from typing import Any, Literal
 
 import torch
 
-from ...core import Chainable, Transform
+from ...core import Chainable, TensorTransform
 from ...utils import NumberList, TensorList, tofloat, unpack_dicts, unpack_states
-from ...utils.linalg.linear_operator import ScaledIdentity
+from ...linalg.linear_operator import ScaledIdentity
 from ..functional import epsilon_step_size
 
 def _acceptable_alpha(alpha, param:torch.Tensor):
@@ -16,7 +16,7 @@ def _acceptable_alpha(alpha, param:torch.Tensor):
         return False
     return True
 
-def _get_H(self: Transform, var):
+def _get_H(self: TensorTransform, var):
     n = sum(p.numel() for p in var.params)
     p = var.params[0]
     alpha = self.global_state.get('alpha', 1)
@@ -25,7 +25,7 @@ def _get_H(self: Transform, var):
     return ScaledIdentity(1 / alpha, shape=(n,n), device=p.device, dtype=p.dtype)
 
 
-class PolyakStepSize(Transform):
+class PolyakStepSize(TensorTransform):
     """Polyak's subgradient method with known or unknown f*.
 
     Args:
@@ -47,7 +47,7 @@ class PolyakStepSize(Transform):
         super().__init__(defaults, uses_grad=use_grad, uses_loss=True, inner=inner)
 
     @torch.no_grad
-    def update_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
         assert grads is not None and loss is not None
         tensors = TensorList(tensors)
         grads = TensorList(grads)
@@ -79,15 +79,15 @@ class PolyakStepSize(Transform):
         self.global_state['alpha'] = alpha
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         alpha = self.global_state.get('alpha', 1)
         if not _acceptable_alpha(alpha, tensors[0]): alpha = epsilon_step_size(TensorList(tensors))
 
         torch._foreach_mul_(tensors, alpha * unpack_dicts(settings, 'alpha', cls=NumberList))
         return tensors
 
-    def get_H(self, var):
-        return _get_H(self, var)
+    def get_H(self, objective):
+        return _get_H(self, objective)
 
 
 def _bb_short(s: TensorList, y: TensorList, sy, eps):
@@ -116,7 +116,7 @@ def _bb_geom(s: TensorList, y: TensorList, sy, eps, fallback:bool):
         return None
     return (short * long) ** 0.5
 
-class BarzilaiBorwein(Transform):
+class BarzilaiBorwein(TensorTransform):
     """Barzilai-Borwein step size method.
 
     Args:
@@ -144,7 +144,7 @@ class BarzilaiBorwein(Transform):
         self.global_state['reset'] = True
 
     @torch.no_grad
-    def update_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
@@ -175,11 +175,11 @@ class BarzilaiBorwein(Transform):
         prev_p.copy_(params)
         prev_g.copy_(g)
 
-    def get_H(self, var):
-        return _get_H(self, var)
+    def get_H(self, objective):
+        return _get_H(self, objective)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         alpha = self.global_state.get('alpha', None)
 
         if not _acceptable_alpha(alpha, tensors[0]):
@@ -189,7 +189,7 @@ class BarzilaiBorwein(Transform):
         return tensors
 
 
-class BBStab(Transform):
+class BBStab(TensorTransform):
     """Stabilized Barzilai-Borwein method (https://arxiv.org/abs/1907.06409).
 
     This clips the norm of the Barzilai-Borwein update by ``delta``, where ``delta`` can be adaptive if ``c`` is specified.
@@ -228,7 +228,7 @@ class BBStab(Transform):
         self.global_state['reset'] = True
 
     @torch.no_grad
-    def update_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
@@ -287,11 +287,11 @@ class BBStab(Transform):
         prev_p.copy_(params)
         prev_g.copy_(g)
 
-    def get_H(self, var):
-        return _get_H(self, var)
+    def get_H(self, objective):
+        return _get_H(self, objective)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         alpha = self.global_state.get('alpha', None)
 
         if not _acceptable_alpha(alpha, tensors[0]):
@@ -301,7 +301,7 @@ class BBStab(Transform):
         return tensors
 
 
-class AdGD(Transform):
+class AdGD(TensorTransform):
     """AdGD and AdGD-2 (https://arxiv.org/abs/2308.02261)"""
     def __init__(self, variant:Literal[1,2]=2, alpha_0:float = 1e-7, sqrt:bool=True, use_grad=True, inner: Chainable | None = None,):
         defaults = dict(variant=variant, alpha_0=alpha_0, sqrt=sqrt)
@@ -313,7 +313,7 @@ class AdGD(Transform):
         self.global_state['reset'] = True
 
     @torch.no_grad
-    def update_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
         variant = settings[0]['variant']
         theta_0 = 0 if variant == 1 else 1/3
         theta = self.global_state.get('theta', theta_0)
@@ -371,7 +371,7 @@ class AdGD(Transform):
         prev_g.copy_(g)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         alpha = self.global_state.get('alpha', None)
 
         if not _acceptable_alpha(alpha, tensors[0]):
@@ -383,5 +383,5 @@ class AdGD(Transform):
         torch._foreach_mul_(tensors, alpha)
         return tensors
 
-    def get_H(self, var):
-        return _get_H(self, var)
+    def get_H(self, objective):
+        return _get_H(self, objective)

torchzero/modules/step_size/lr.py

@@ -2,7 +2,7 @@
 import torch
 import random
 
-from ...core import Transform
+from ...core import TensorTransform
 from ...utils import NumberList, TensorList, generic_ne, unpack_dicts
 
 def lazy_lr(tensors: TensorList, lr: float | list, inplace:bool):
@@ -12,24 +12,24 @@ def lazy_lr(tensors: TensorList, lr: float | list, inplace:bool):
         return tensors * lr
     return tensors
 
-class LR(Transform):
+class LR(TensorTransform):
     """Learning rate. Adding this module also adds support for LR schedulers."""
     def __init__(self, lr: float):
         defaults=dict(lr=lr)
-        super().__init__(defaults, uses_grad=False)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         return lazy_lr(TensorList(tensors), lr=[s['lr'] for s in settings], inplace=True)
 
-class StepSize(Transform):
+class StepSize(TensorTransform):
     """this is exactly the same as LR, except the `lr` parameter can be renamed to any other name to avoid clashes"""
     def __init__(self, step_size: float, key = 'step_size'):
         defaults={"key": key, key: step_size}
-        super().__init__(defaults, uses_grad=False)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         return lazy_lr(TensorList(tensors), lr=[s[s['key']] for s in settings], inplace=True)
 
 
@@ -38,8 +38,8 @@ def _warmup_lr(step: int, start_lr: float | NumberList, end_lr: float | NumberLi
     if step > steps: return end_lr
     return start_lr + (end_lr - start_lr) * (step / steps)
 
-class Warmup(Transform):
-    """Learning rate warmup, linearly increases learning rate multiplier from :code:`start_lr` to :code:`end_lr` over :code:`steps` steps.
+class Warmup(TensorTransform):
+    """Learning rate warmup, linearly increases learning rate multiplier from ``start_lr`` to ``end_lr`` over ``steps`` steps.
 
     Args:
         steps (int, optional): number of steps to perform warmup for. Defaults to 100.
@@ -64,7 +64,7 @@ class Warmup(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         start_lr, end_lr = unpack_dicts(settings, 'start_lr', 'end_lr', cls = NumberList)
         num_steps = settings[0]['steps']
         step = self.global_state.get('step', 0)
@@ -77,7 +77,7 @@ class Warmup(Transform):
         self.global_state['step'] = step + 1
         return tensors
 
-class WarmupNormClip(Transform):
+class WarmupNormClip(TensorTransform):
     """Warmup via clipping of the update norm.
 
     Args:
@@ -102,7 +102,7 @@ class WarmupNormClip(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         start_norm, end_norm = unpack_dicts(settings, 'start_norm', 'end_norm', cls = NumberList)
         num_steps = settings[0]['steps']
         step = self.global_state.get('step', 0)
@@ -118,8 +118,8 @@ class WarmupNormClip(Transform):
         return tensors
 
 
-class RandomStepSize(Transform):
-    """Uses random global or layer-wise step size from `low` to `high`.
+class RandomStepSize(TensorTransform):
+    """Uses random global or layer-wise step size from ``low`` to ``high``.
 
     Args:
         low (float, optional): minimum learning rate. Defaults to 0.
@@ -133,7 +133,7 @@ class RandomStepSize(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         s = settings[0]
         parameterwise = s['parameterwise']
 

torchzero/modules/termination/termination.py

@@ -1,11 +1,11 @@
 import time
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
-from typing import cast
+from typing import cast, final
 
 import torch
 
-from ...core import Module, Var
+from ...core import Module, Objective
 from ...utils import Metrics, TensorList, safe_dict_update_, tofloat
 
 
@@ -16,14 +16,15 @@ class TerminationCriteriaBase(Module):
         super().__init__(defaults)
 
     @abstractmethod
-    def termination_criteria(self, var: Var) -> bool:
+    def termination_criteria(self, objective: Objective) -> bool:
         ...
 
-    def should_terminate(self, var: Var) -> bool:
+    @final
+    def should_terminate(self, objective: Objective) -> bool:
         n_bad = self.global_state.get('_n_bad', 0)
         n = self.defaults['_n']
 
-        if self.termination_criteria(var):
+        if self.termination_criteria(objective):
             n_bad += 1
             if n_bad >= n:
                 self.global_state['_n_bad'] = 0
@@ -36,12 +37,12 @@ class TerminationCriteriaBase(Module):
         return False
 
 
-    def update(self, var):
-        var.should_terminate = self.should_terminate(var)
-        if var.should_terminate: self.global_state['_n_bad'] = 0
+    def update(self, objective):
+        objective.should_terminate = self.should_terminate(objective)
+        if objective.should_terminate: self.global_state['_n_bad'] = 0
 
-    def apply(self, var):
-        return var
+    def apply(self, objective):
+        return objective
 
 
 class TerminateAfterNSteps(TerminationCriteriaBase):
@@ -49,7 +50,7 @@ class TerminateAfterNSteps(TerminationCriteriaBase):
         defaults = dict(steps=steps)
         super().__init__(defaults)
 
-    def termination_criteria(self, var):
+    def termination_criteria(self, objective):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
@@ -61,16 +62,17 @@ class TerminateAfterNEvaluations(TerminationCriteriaBase):
         defaults = dict(maxevals=maxevals)
         super().__init__(defaults)
 
-    def termination_criteria(self, var):
+    def termination_criteria(self, objective):
         maxevals = self.defaults['maxevals']
-        return var.modular.num_evaluations >= maxevals
+        assert objective.modular is not None
+        return objective.modular.num_evaluations >= maxevals
 
 class TerminateAfterNSeconds(TerminationCriteriaBase):
     def __init__(self, seconds:float, sec_fn = time.time):
         defaults = dict(seconds=seconds, sec_fn=sec_fn)
         super().__init__(defaults)
 
-    def termination_criteria(self, var):
+    def termination_criteria(self, objective):
         max_seconds = self.defaults['seconds']
         sec_fn = self.defaults['sec_fn']
 
@@ -88,10 +90,10 @@ class TerminateByGradientNorm(TerminationCriteriaBase):
         defaults = dict(tol=tol, ord=ord)
         super().__init__(defaults, n=n)
 
-    def termination_criteria(self, var):
+    def termination_criteria(self, objective):
         tol = self.defaults['tol']
         ord = self.defaults['ord']
-        return TensorList(var.get_grad()).global_metric(ord) <= tol
+        return TensorList(objective.get_grads()).global_metric(ord) <= tol
 
 
 class TerminateByUpdateNorm(TerminationCriteriaBase):
@@ -100,20 +102,20 @@ class TerminateByUpdateNorm(TerminationCriteriaBase):
         defaults = dict(tol=tol, ord=ord)
         super().__init__(defaults, n=n)
 
-    def termination_criteria(self, var):
+    def termination_criteria(self, objective):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
         tol = self.defaults['tol']
         ord = self.defaults['ord']
 
-        p_prev = self.get_state(var.params, 'p_prev', cls=TensorList)
+        p_prev = self.get_state(objective.params, 'p_prev', cls=TensorList)
         if step == 0:
-            p_prev.copy_(var.params)
+            p_prev.copy_(objective.params)
             return False
 
-        should_terminate = (p_prev - var.params).global_metric(ord) <= tol
-        p_prev.copy_(var.params)
+        should_terminate = (p_prev - objective.params).global_metric(ord) <= tol
+        p_prev.copy_(objective.params)
         return should_terminate
 
 
@@ -122,10 +124,10 @@ class TerminateOnNoImprovement(TerminationCriteriaBase):
         defaults = dict(tol=tol)
         super().__init__(defaults, n=n)
 
-    def termination_criteria(self, var):
+    def termination_criteria(self, objective):
         tol = self.defaults['tol']
 
-        f = tofloat(var.get_loss(False))
+        f = tofloat(objective.get_loss(False))
         if 'f_min' not in self.global_state:
             self.global_state['f_min'] = f
             return False
@@ -141,9 +143,9 @@ class TerminateOnLossReached(TerminationCriteriaBase):
         defaults = dict(value=value)
         super().__init__(defaults)
 
-    def termination_criteria(self, var):
+    def termination_criteria(self, objective):
         value = self.defaults['value']
-        return var.get_loss(False) <= value
+        return objective.get_loss(False) <= value
 
 class TerminateAny(TerminationCriteriaBase):
     def __init__(self, *criteria: TerminationCriteriaBase):
@@ -151,9 +153,9 @@ class TerminateAny(TerminationCriteriaBase):
 
         self.set_children_sequence(criteria)
 
-    def termination_criteria(self, var: Var) -> bool:
+    def termination_criteria(self, objective: Objective) -> bool:
         for c in self.get_children_sequence():
-            if cast(TerminationCriteriaBase, c).termination_criteria(var): return True
+            if cast(TerminationCriteriaBase, c).termination_criteria(objective): return True
 
         return False
 
@@ -163,9 +165,9 @@ class TerminateAll(TerminationCriteriaBase):
 
         self.set_children_sequence(criteria)
 
-    def termination_criteria(self, var: Var) -> bool:
+    def termination_criteria(self, objective: Objective) -> bool:
         for c in self.get_children_sequence():
-            if not cast(TerminationCriteriaBase, c).termination_criteria(var): return False
+            if not cast(TerminationCriteriaBase, c).termination_criteria(objective): return False
 
         return True
 
@@ -173,7 +175,7 @@ class TerminateNever(TerminationCriteriaBase):
     def __init__(self):
         super().__init__()
 
-    def termination_criteria(self, var): return False
+    def termination_criteria(self, objective): return False
 
 def make_termination_criteria(
     ftol: float | None = None,

torchzero/modules/trust_region/cubic_regularization.py

@@ -5,7 +5,7 @@ import torch
 
 from ...core import Chainable, Module
 from ...utils import TensorList, vec_to_tensors
-from ...utils.linalg.linear_operator import LinearOperator
+from ...linalg.linear_operator import LinearOperator
 from .trust_region import _RADIUS_KEYS, TrustRegionBase, _RadiusStrategy
 
 
@@ -58,7 +58,7 @@ def ls_cubic_solver(f, g:torch.Tensor, H:LinearOperator, M: float, loss_at_param
     for _ in range(it_max):
         r_try = (r_min + r_max) / 2
         lam = r_try * M
-        s_lam = H.add_diagonal(lam).solve(g).neg()
+        s_lam = H.solve_plus_diag(g, lam).neg()
         # s_lam = -torch.linalg.solve(B + lam*id_matrix, g)
         solver_it += 1
         crit = conv_criterion(s_lam, r_try)

torchzero/modules/trust_region/levenberg_marquardt.py

@@ -2,7 +2,7 @@
 import torch
 
 from ...core import Chainable, Module
-from ...utils.linalg import linear_operator
+from ...linalg import linear_operator
 from .trust_region import _RADIUS_KEYS, TrustRegionBase, _RadiusStrategy
 
 
@@ -32,38 +32,31 @@ class LevenbergMarquardt(TrustRegionBase):
         max_attempts (max_attempts, optional):
             maximum number of trust region size size reductions per step. A zero update vector is returned when
             this limit is exceeded. Defaults to 10.
+        adaptive (bool, optional):
+            if True, trust radius is multiplied by square root of gradient norm.
         fallback (bool, optional):
             if ``True``, when ``hess_module`` maintains hessian inverse which can't be inverted efficiently, it will
             be inverted anyway. When ``False`` (default), a ``RuntimeError`` will be raised instead.
         inner (Chainable | None, optional): preconditioning is applied to output of thise module. Defaults to None.
 
-    Examples:
-        Gauss-Newton with Levenberg-Marquardt trust-region
+    ### Examples:
 
-        .. code-block:: python
+    Gauss-Newton with Levenberg-Marquardt trust-region
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.LevenbergMarquardt(tz.m.GaussNewton()),
-            )
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.LevenbergMarquardt(tz.m.GaussNewton()),
+    )
+    ```
 
-        LM-SR1
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.LevenbergMarquardt(tz.m.SR1(inverse=False)),
-            )
-
-        First order trust region (hessian is assumed to be identity)
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.LevenbergMarquardt(tz.m.Identity()),
-            )
+    LM-SR1
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.LevenbergMarquardt(tz.m.SR1(inverse=False)),
+    )
+    ```
 
     """
     def __init__(
@@ -78,11 +71,12 @@ class LevenbergMarquardt(TrustRegionBase):
         max_attempts: int = 10,
         radius_strategy: _RadiusStrategy | _RADIUS_KEYS = 'default',
         y: float = 0,
+        adaptive: bool = False,
         fallback: bool = False,
         update_freq: int = 1,
         inner: Chainable | None = None,
     ):
-        defaults = dict(y=y, fallback=fallback)
+        defaults = dict(y=y, fallback=fallback, adaptive=adaptive)
         super().__init__(
             defaults=defaults,
             hess_module=hess_module,
@@ -103,6 +97,7 @@ class LevenbergMarquardt(TrustRegionBase):
 
     def trust_solve(self, f, g, H, radius, params, closure, settings):
         y = settings['y']
+        adaptive = settings["adaptive"]
 
         if isinstance(H, linear_operator.DenseInverse):
             if settings['fallback']:
@@ -117,12 +112,14 @@ class LevenbergMarquardt(TrustRegionBase):
                 )
 
         reg = 1/radius
+        if adaptive: reg = reg * torch.linalg.vector_norm(g).sqrt()
+
         if y == 0:
-            return H.add_diagonal(reg).solve(g)
+            return H.solve_plus_diag(g, reg) # pyright:ignore[reportAttributeAccessIssue]
 
         diag = H.diagonal()
         diag = torch.where(diag < torch.finfo(diag.dtype).tiny * 2, 1, diag)
         if y != 1: diag = (diag*y) + (1-y)
-        return H.add_diagonal(diag*reg).solve(g)
+        return H.solve_plus_diag(g, diag*reg)