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

torchzero/modules/optimizers/sophia_h.py

@@ -1,129 +0,0 @@
-from typing import Literal
-from collections.abc import Callable
-import torch
-
-from ...core import Module, Target, Transform, Chainable, apply_transform
-from ...utils import NumberList, TensorList, as_tensorlist
-from ...utils.derivatives import hvp, hvp_fd_forward, hvp_fd_central
-
-def sophia_H(
-    tensors: TensorList,
-    h: TensorList | None,
-    exp_avg_: TensorList,
-    h_exp_avg_: TensorList,
-    beta1: float | NumberList,
-    beta2: float | NumberList,
-    update_freq: int,
-    precond_scale: float | NumberList,
-    clip: float | NumberList,
-    eps: float | NumberList,
-    step: int
-):
-    # momentum
-    exp_avg_.lerp_(tensors, 1-beta1)
-
-    # update preconditioner
-    if step % update_freq == 0:
-        assert h is not None
-        h_exp_avg_.lerp_(h, 1-beta2)
-
-    else:
-        assert h is None
-
-    denom = (h_exp_avg_ * precond_scale).clip_(min=eps)
-    return (exp_avg_ / denom).clip_(-clip, clip)
-
-
-class SophiaH(Module):
-    def __init__(
-        self,
-        beta1: float = 0.96,
-        beta2: float = 0.99,
-        update_freq: int = 10,
-        precond_scale: float = 1,
-        clip: float = 1,
-        eps: float = 1e-12,
-        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, precond_scale=precond_scale, clip=clip, 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, precond_scale, clip, eps = self.get_settings(params,
-            'beta1', 'beta2', 'precond_scale', 'clip', 'eps', cls=NumberList)
-
-        exp_avg, h_exp_avg = 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
-
-        h = None
-        if step % update_freq == 0:
-
-            grad=None
-            for i in range(n_samples):
-                u = [torch.randn(p.shape, device=p.device, dtype=p.dtype, generator=generator) for p in params]
-
-                if hvp_method == 'autograd':
-                    if grad is None: grad = var.get_grad(create_graph=True)
-                    assert grad is not None
-                    Hvp = hvp(params, grad, u, retain_graph=i < n_samples-1)
-
-                elif hvp_method == 'forward':
-                    loss, Hvp = hvp_fd_forward(closure, params, u, h=fd_h, g_0=var.get_grad(), normalize=True)
-
-                elif hvp_method == 'central':
-                    loss, Hvp = hvp_fd_central(closure, params, u, h=fd_h, normalize=True)
-
-                else:
-                    raise ValueError(hvp_method)
-
-                if h is None: h = Hvp
-                else: torch._foreach_add_(h, Hvp)
-
-            assert h is not None
-            if n_samples > 1: torch._foreach_div_(h, n_samples)
-
-        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 = sophia_H(
-            tensors=TensorList(update),
-            h=TensorList(h) if h is not None else None,
-            exp_avg_=exp_avg,
-            h_exp_avg_=h_exp_avg,
-            beta1=beta1,
-            beta2=beta2,
-            update_freq=update_freq,
-            precond_scale=precond_scale,
-            clip=clip,
-            eps=eps,
-            step=step,
-        )
-        return var

torchzero/modules/quasi_newton/cg.py

@@ -1,268 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Literal
-
-import torch
-
-from ...core import Chainable, TensorwiseTransform, Transform, apply_transform
-from ...utils import TensorList, as_tensorlist, unpack_dicts, unpack_states
-
-
-class ConguateGradientBase(Transform, ABC):
-    """all CGs are the same except beta calculation"""
-    def __init__(self, defaults = None, clip_beta: bool = False, reset_interval: int | None | Literal['auto'] = None, inner: Chainable | None = None):
-        if defaults is None: defaults = {}
-        defaults['reset_interval'] = reset_interval
-        defaults['clip_beta'] = clip_beta
-        super().__init__(defaults, uses_grad=False)
-
-        if inner is not None:
-            self.set_child('inner', inner)
-
-    def initialize(self, p: TensorList, g: TensorList):
-        """runs on first step when prev_grads and prev_dir are not available"""
-
-    @abstractmethod
-    def get_beta(self, p: TensorList, g: TensorList, prev_g: TensorList, prev_d: TensorList) -> float | torch.Tensor:
-        """returns beta"""
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        tensors = as_tensorlist(tensors)
-        params = as_tensorlist(params)
-
-        step = self.global_state.get('step', 0)
-        prev_dir, prev_grads = unpack_states(states, tensors, 'prev_dir', 'prev_grad', cls=TensorList)
-
-        # initialize on first step
-        if step == 0:
-            self.initialize(params, tensors)
-            prev_dir.copy_(tensors)
-            prev_grads.copy_(tensors)
-            self.global_state['step'] = step + 1
-            return tensors
-
-        # get beta
-        beta = self.get_beta(params, tensors, prev_grads, prev_dir)
-        if settings[0]['clip_beta']: beta = max(0, beta) # pyright:ignore[reportArgumentType]
-        prev_grads.copy_(tensors)
-
-        # inner step
-        if 'inner' in self.children:
-            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads))
-
-        # calculate new direction with beta
-        dir = tensors.add_(prev_dir.mul_(beta))
-        prev_dir.copy_(dir)
-
-        # resetting
-        self.global_state['step'] = step + 1
-        reset_interval = settings[0]['reset_interval']
-        if reset_interval == 'auto': reset_interval = tensors.global_numel() + 1
-        if reset_interval is not None and (step+1) % reset_interval == 0:
-            self.reset()
-
-        return dir
-
-# ------------------------------- Polak-Ribière ------------------------------ #
-def polak_ribiere_beta(g: TensorList, prev_g: TensorList):
-    denom = prev_g.dot(prev_g)
-    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
-    return g.dot(g - prev_g) / denom
-
-class PolakRibiere(ConguateGradientBase):
-    """Polak-Ribière-Polyak nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe(c2=0.1)` after this."""
-    def __init__(self, clip_beta=True, reset_interval: int | None = None, inner: Chainable | None = None):
-        super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
-
-    def get_beta(self, p, g, prev_g, prev_d):
-        return polak_ribiere_beta(g, prev_g)
-
-# ------------------------------ Fletcher–Reeves ----------------------------- #
-def fletcher_reeves_beta(gg: torch.Tensor, prev_gg: torch.Tensor):
-    if prev_gg.abs() <= torch.finfo(gg.dtype).eps: return 0
-    return gg / prev_gg
-
-class FletcherReeves(ConguateGradientBase):
-    """Fletcher–Reeves nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None | Literal['auto'] = 'auto', clip_beta=False, inner: Chainable | None = None):
-        super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
-
-    def initialize(self, p, g):
-        self.global_state['prev_gg'] = g.dot(g)
-
-    def get_beta(self, p, g, prev_g, prev_d):
-        gg = g.dot(g)
-        beta = fletcher_reeves_beta(gg, self.global_state['prev_gg'])
-        self.global_state['prev_gg'] = gg
-        return beta
-
-# ----------------------------- Hestenes–Stiefel ----------------------------- #
-def hestenes_stiefel_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
-    grad_diff = g - prev_g
-    denom = prev_d.dot(grad_diff)
-    if denom.abs() < torch.finfo(g[0].dtype).eps: return 0
-    return (g.dot(grad_diff) / denom).neg()
-
-
-class HestenesStiefel(ConguateGradientBase):
-    """Hestenes–Stiefel nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
-        super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
-
-    def get_beta(self, p, g, prev_g, prev_d):
-        return hestenes_stiefel_beta(g, prev_d, prev_g)
-
-
-# --------------------------------- Dai–Yuan --------------------------------- #
-def dai_yuan_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
-    denom = prev_d.dot(g - prev_g)
-    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
-    return (g.dot(g) / denom).neg()
-
-class DaiYuan(ConguateGradientBase):
-    """Dai–Yuan nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
-        super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
-
-    def get_beta(self, p, g, prev_g, prev_d):
-        return dai_yuan_beta(g, prev_d, prev_g)
-
-
-# -------------------------------- Liu-Storey -------------------------------- #
-def liu_storey_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList, ):
-    denom = prev_g.dot(prev_d)
-    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
-    return g.dot(g - prev_g) / denom
-
-class LiuStorey(ConguateGradientBase):
-    """Liu-Storey nonlinear conjugate gradient method. This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
-        super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
-
-    def get_beta(self, p, g, prev_g, prev_d):
-        return liu_storey_beta(g, prev_d, prev_g)
-
-# ----------------------------- Conjugate Descent ---------------------------- #
-class ConjugateDescent(Transform):
-    """Conjugate Descent (CD). This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, inner: Chainable | None = None):
-        super().__init__(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):
-        g = as_tensorlist(tensors)
-
-        prev_d = unpack_states(states, tensors, 'prev_dir', cls=TensorList, init=torch.zeros_like)
-        if 'denom' not in self.global_state:
-            self.global_state['denom'] = torch.tensor(0.).to(g[0])
-
-        prev_gd = self.global_state.get('prev_gd', 0)
-        if abs(prev_gd) <= torch.finfo(g[0].dtype).eps: beta = 0
-        else: beta = g.dot(g) / prev_gd
-
-        # inner step
-        if 'inner' in self.children:
-            g = as_tensorlist(apply_transform(self.children['inner'], g, params, grads))
-
-        dir = g.add_(prev_d.mul_(beta))
-        prev_d.copy_(dir)
-        self.global_state['prev_gd'] = g.dot(dir)
-        return dir
-
-
-# -------------------------------- Hager-Zhang ------------------------------- #
-def hager_zhang_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList,):
-    g_diff = g - prev_g
-    denom = prev_d.dot(g_diff)
-    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
-
-    term1 = 1/denom
-    # term2
-    term2 = (g_diff - (2 * prev_d * (g_diff.pow(2).global_sum()/denom))).dot(g)
-    return (term1 * term2).neg()
-
-
-class HagerZhang(ConguateGradientBase):
-    """Hager-Zhang nonlinear conjugate gradient method,
-    This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
-        super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
-
-    def get_beta(self, p, g, prev_g, prev_d):
-        return hager_zhang_beta(g, prev_d, prev_g)
-
-
-# ----------------------------------- HS-DY ---------------------------------- #
-def hs_dy_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
-    grad_diff = g - prev_g
-    denom = prev_d.dot(grad_diff)
-    if denom.abs() <= torch.finfo(g[0].dtype).eps: return 0
-
-    # Dai-Yuan
-    dy_beta = (g.dot(g) / denom).neg().clamp(min=0)
-
-    # Hestenes–Stiefel
-    hs_beta = (g.dot(grad_diff) / denom).neg().clamp(min=0)
-
-    return max(0, min(dy_beta, hs_beta)) # type:ignore
-
-class HybridHS_DY(ConguateGradientBase):
-    """HS-DY hybrid conjugate gradient method.
-    This requires step size to be determined via a line search, so put a line search like :code:`StrongWolfe` after this."""
-    def __init__(self, reset_interval: int | None | Literal['auto'] = None, clip_beta=False, inner: Chainable | None = None):
-        super().__init__(clip_beta=clip_beta, reset_interval=reset_interval, inner=inner)
-
-    def get_beta(self, p, g, prev_g, prev_d):
-        return hs_dy_beta(g, prev_d, prev_g)
-
-
-def projected_gradient_(H:torch.Tensor, y:torch.Tensor, tol: float):
-    Hy = H @ y
-    denom = y.dot(Hy)
-    if denom.abs() < tol: return H
-    H -= (H @ y.outer(y) @ H) / denom
-    return H
-
-class ProjectedGradientMethod(TensorwiseTransform):
-    """Pearson, J. D. (1969). Variable metric methods of minimisation. The Computer Journal, 12(2), 171–178. doi:10.1093/comjnl/12.2.171.
-
-    (This is not the same as projected gradient descent)
-    """
-
-    def __init__(
-        self,
-        tol: float = 1e-10,
-        reset_interval: int | None = None,
-        update_freq: int = 1,
-        scale_first: bool = False,
-        concat_params: bool = True,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(reset_interval=reset_interval, tol=tol)
-        super().__init__(defaults, uses_grad=False, scale_first=scale_first, concat_params=concat_params, update_freq=update_freq, inner=inner)
-
-    def update_tensor(self, tensor, param, grad, loss, state, settings):
-        step = state.get('step', 0)
-        state['step'] = step + 1
-        reset_interval = settings['reset_interval']
-        if reset_interval is None: reset_interval = tensor.numel() + 1 # as recommended
-
-        if ("H" not in state) or (step % reset_interval == 0):
-            state["H"] = torch.eye(tensor.numel(), device=tensor.device, dtype=tensor.dtype)
-            state['g_prev'] = tensor.clone()
-            return
-
-        H = state['H']
-        g_prev = state['g_prev']
-        state['g_prev'] = tensor.clone()
-        y = (tensor - g_prev).ravel()
-
-        projected_gradient_(H, y, settings['tol'])
-
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
-        H = state['H']
-        return (H @ tensor.view(-1)).view_as(tensor)

torchzero/modules/quasi_newton/experimental/__init__.py

@@ -1 +0,0 @@
-from .modular_lbfgs import ModularLBFGS

torchzero/modules/quasi_newton/experimental/modular_lbfgs.py

@@ -1,266 +0,0 @@
-from collections import deque
-from operator import itemgetter
-from typing import Any
-
-import torch
-
-from ....core import Chainable, Module, Transform, Var, apply_transform, maybe_chain
-from ....utils import NumberList, TensorList, as_tensorlist
-
-
-def _adaptive_damping(
-    s_k: TensorList,
-    y_k: TensorList,
-    ys_k: torch.Tensor,
-    init_damping = 0.99,
-    eigval_bounds = (0.01, 1.5)
-):
-    # adaptive damping Al-Baali, M.: Quasi-Wolfe conditions for quasi-Newton methods for large-scale optimization. In: 40th Workshop on Large Scale Nonlinear Optimization, Erice, Italy, June 22–July 1 (2004)
-    sigma_l, sigma_h = eigval_bounds
-    u = ys_k / s_k.dot(s_k)
-    if u <= sigma_l < 1: tau = min((1-sigma_l)/(1-u), init_damping)
-    elif u >= sigma_h > 1: tau = min((sigma_h-1)/(u-1), init_damping)
-    else: tau = init_damping
-    y_k = tau * y_k + (1-tau) * s_k
-    ys_k = s_k.dot(y_k)
-
-    return s_k, y_k, ys_k
-
-def lbfgs(
-    tensors_: TensorList,
-    var: Var,
-    s_history: deque[TensorList],
-    y_history: deque[TensorList],
-    sy_history: deque[torch.Tensor],
-    y_k: TensorList | None,
-    ys_k: torch.Tensor | None,
-    z_tfm: Any,
-):
-    if len(s_history) == 0 or y_k is None or ys_k is None:
-
-        # initial step size guess modified from pytorch L-BFGS
-        scale = 1 / tensors_.abs().global_sum()
-        if scale < 1e-5: scale = 1 / tensors_.abs().mean()
-        return tensors_.mul_(min(1.0, scale)) # pyright: ignore[reportArgumentType]
-
-    else:
-        # 1st loop
-        alpha_list = []
-        q = tensors_.clone()
-        for s_i, y_i, ys_i in zip(reversed(s_history), reversed(y_history), reversed(sy_history)):
-            p_i = 1 / ys_i # this is also denoted as ρ (rho)
-            alpha = p_i * s_i.dot(q)
-            alpha_list.append(alpha)
-            q.sub_(y_i, alpha=alpha) # pyright: ignore[reportArgumentType]
-
-        # calculate z
-        # s.y/y.y is also this weird y-looking symbol I couldn't find
-        # z is it times q
-        # actually H0 = (s.y/y.y) * I, and z = H0 @ q
-        z = q * (ys_k / (y_k.dot(y_k)))
-
-        if z_tfm is not None:
-            z = TensorList(apply_transform(z_tfm, tensors=z, params=var.params, grads=var.grad, var=var))
-
-        # 2nd loop
-        for s_i, y_i, ys_i, alpha_i in zip(s_history, y_history, sy_history, reversed(alpha_list)):
-            p_i = 1 / ys_i
-            beta_i = p_i * y_i.dot(z)
-            z.add_(s_i, alpha = alpha_i - beta_i)
-
-        return z
-
-def _apply_tfms_into_history(
-    self: Module,
-    params: list[torch.Tensor],
-    var: Var,
-    update: list[torch.Tensor],
-):
-    if 'params_history_tfm' in self.children:
-        params = apply_transform(self.children['params_history_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=var.grad, var=var)
-
-    if 'grad_history_tfm' in self.children:
-        update = apply_transform(self.children['grad_history_tfm'], tensors=as_tensorlist(update).clone(), params=params, grads=var.grad, var=var)
-
-    return params, update
-
-def _apply_tfms_into_precond(
-    self: Module,
-    params: list[torch.Tensor],
-    var: Var,
-    update: list[torch.Tensor],
-):
-    if 'params_precond_tfm' in self.children:
-        params = apply_transform(self.children['params_precond_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=var.grad, var=var)
-
-    if 'grad_precond_tfm' in self.children:
-        update = apply_transform(self.children['grad_precond_tfm'], tensors=update, params=params, grads=var.grad, var=var)
-
-    return params, update
-
-
-class ModularLBFGS(Module):
-    """L-BFGS with ability to apply transforms to many inner variables.
-
-    Args:
-        history_size (int, optional): number of past parameter differences and gradient differences to store. Defaults to 10.
-        tol (float | None, optional):
-            tolerance for minimal gradient difference to avoid instability after converging to minima. Defaults to 1e-10.
-        damping (bool, optional):
-            whether to use adaptive damping. Learning rate might need to be lowered with this enabled. Defaults to False.
-        init_damping (float, optional):
-            initial damping for adaptive dampening. Defaults to 0.9.
-        eigval_bounds (tuple, optional):
-            eigenvalue bounds for adaptive dampening. Defaults to (0.5, 50).
-        update_freq (int, optional):
-            how often to update L-BFGS history. Defaults to 1.
-        z_tfm (float | None, optional):
-            transform module applied to initial H^-1 @ q guess. Defaults to None.
-        params_history_tfm (AnyTransform | None, optional):
-            transform module applied to params before adding s_k to history. Defaults to None.
-        grad_history_tfm (AnyTransform | None, optional):
-            transform module applied to grads before adding y_k to history. Defaults to None.
-        params_precond_tfm (AnyTransform | None, optional):
-            transform module applied to params to calculate s_k before preconditioning. Defaults to None.
-        grad_precond_tfm (AnyTransform | None, optional):
-            transform module applied to grads to calculate y_k before preconditioning. Defaults to None.
-        update_precond_tfm (Chainable | None, optional):
-            transform module applied to grads that are being preconditioned. Defaults to None.
-    """
-    def __init__(
-        self,
-        history_size=10,
-        tol: float | None = 1e-10,
-        damping: bool = False,
-        init_damping=0.9,
-        eigval_bounds=(0.5, 50),
-        update_freq = 1,
-        params_history_tfm: Chainable | None = None,
-        grad_history_tfm: Chainable | None = None,
-        params_precond_tfm: Chainable | None = None,
-        grad_precond_tfm: Chainable | None = None,
-        update_precond_tfm: Chainable | None = None,
-        z_tfm: Chainable | None = None,
-    ):
-        defaults = dict(history_size=history_size, tol=tol, damping=damping, init_damping=init_damping, eigval_bounds=eigval_bounds, update_freq=update_freq)
-        super().__init__(defaults)
-
-        self.global_state['s_history'] = deque(maxlen=history_size)
-        self.global_state['y_history'] = deque(maxlen=history_size)
-        self.global_state['sy_history'] = deque(maxlen=history_size)
-
-        loc = locals().copy()
-        for k in ('update_precond_tfm', 'params_history_tfm', 'grad_history_tfm', 'params_precond_tfm', 'grad_precond_tfm','z_tfm'):
-            v = loc[k]
-            if v is not None:
-                self.set_child(k,v)
-
-    def reset(self):
-        """Resets the internal state of the L-SR1 module."""
-        # super().reset() # Clears self.state (per-parameter) if any, and "step"
-        self.state.clear()
-        self.global_state['step'] = 0
-        self.global_state['s_history'].clear()
-        self.global_state['y_history'].clear()
-        self.global_state['sy_history'].clear()
-
-    @torch.no_grad
-    def step(self, var):
-        params = as_tensorlist(var.params)
-        update = as_tensorlist(var.get_update())
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
-
-        # history of s and k
-        s_history: deque[TensorList] = self.global_state['s_history']
-        y_history: deque[TensorList] = self.global_state['y_history']
-        sy_history: deque[torch.Tensor] = self.global_state['sy_history']
-
-        tol, damping, init_damping, eigval_bounds, update_freq = itemgetter(
-            'tol', 'damping', 'init_damping', 'eigval_bounds', 'update_freq')(self.settings[params[0]])
-
-        # params_beta, grads_beta = self.get_settings('params_beta', 'grads_beta', params=params, cls=NumberList)
-        # l_params, l_update = _lerp_params_update_(self, params, update, params_beta, grads_beta)
-
-        # params and update that go into history
-        params_h, update_h = _apply_tfms_into_history(
-            self,
-            params=params,
-            var=var,
-            update=update,
-        )
-
-        prev_params_h, prev_grad_h = self.get_state(params, 'prev_params_h', 'prev_grad_h', cls=TensorList)
-
-        # 1st step - there are no previous params and grads, `lbfgs` will do normalized SGD step
-        if step == 0:
-            s_k_h = None; y_k_h = None; ys_k_h = None
-        else:
-            s_k_h = params_h - prev_params_h
-            y_k_h = update_h - prev_grad_h
-            ys_k_h = s_k_h.dot(y_k_h)
-
-            if damping:
-                s_k_h, y_k_h, ys_k_h = _adaptive_damping(s_k_h, y_k_h, ys_k_h, init_damping=init_damping, eigval_bounds=eigval_bounds)
-
-        prev_params_h.copy_(params_h)
-        prev_grad_h.copy_(update_h)
-
-        # update effective preconditioning state
-        if step % update_freq == 0:
-            if ys_k_h is not None and ys_k_h > 1e-10:
-                assert s_k_h is not None and y_k_h is not None
-                s_history.append(s_k_h)
-                y_history.append(y_k_h)
-                sy_history.append(ys_k_h)
-
-        # step with inner module before applying preconditioner
-        if 'update_precond_tfm' in self.children:
-            update_precond_tfm = self.children['update_precond_tfm']
-            inner_var = update_precond_tfm.step(var.clone(clone_update=True))
-            var.update_attrs_from_clone_(inner_var)
-            tensors = inner_var.update
-            assert tensors is not None
-        else:
-            tensors = update.clone()
-
-        # transforms into preconditioner
-        params_p, update_p = _apply_tfms_into_precond(self, params=params, var=var, update=update)
-        prev_params_p, prev_grad_p = self.get_state(params, 'prev_params_p', 'prev_grad_p', cls=TensorList)
-
-        if step == 0:
-            s_k_p = None; y_k_p = None; ys_k_p = None
-
-        else:
-            s_k_p = params_p - prev_params_p
-            y_k_p = update_p - prev_grad_p
-            ys_k_p = s_k_p.dot(y_k_p)
-
-            if damping:
-                s_k_p, y_k_p, ys_k_p = _adaptive_damping(s_k_p, y_k_p, ys_k_p, init_damping=init_damping, eigval_bounds=eigval_bounds)
-
-        prev_params_p.copy_(params_p)
-        prev_grad_p.copy_(update_p)
-
-        # tolerance on gradient difference to avoid exploding after converging
-        if tol is not None:
-            if y_k_p is not None and y_k_p.abs().global_max() <= tol:
-                var.update = update # may have been updated by inner module, probably makes sense to use it here?
-                return var
-
-        # precondition
-        dir = lbfgs(
-            tensors_=as_tensorlist(tensors),
-            var=var,
-            s_history=s_history,
-            y_history=y_history,
-            sy_history=sy_history,
-            y_k=y_k_p,
-            ys_k=ys_k_p,
-            z_tfm=self.children.get('z_tfm', None),
-        )
-
-        var.update = dir
-
-        return var
-