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

torchzero/modules/quasi_newton/cg.py

@@ -1,14 +1,15 @@
 from abc import ABC, abstractmethod
+from typing import Literal
 
 import torch
 
-from ...core import Chainable, Transform, apply
-from ...utils import TensorList, as_tensorlist
+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 = None, inner: Chainable | None = None):
+    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
@@ -25,12 +26,12 @@ class ConguateGradientBase(Transform, ABC):
         """returns beta"""
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    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 = self.get_state('prev_dir', 'prev_grad', params=params, cls=TensorList)
+        prev_dir, prev_grads = unpack_states(states, tensors, 'prev_dir', 'prev_grad', cls=TensorList)
 
         # initialize on first step
         if step == 0:
@@ -42,12 +43,12 @@ class ConguateGradientBase(Transform, ABC):
 
         # get beta
         beta = self.get_beta(params, tensors, prev_grads, prev_dir)
-        if self.settings[params[0]]['clip_beta']: beta = max(0, beta) # pyright:ignore[reportArgumentType]
+        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(self.children['inner'], tensors, params, grads, vars))
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads))
 
         # calculate new direction with beta
         dir = tensors.add_(prev_dir.mul_(beta))
@@ -55,7 +56,8 @@ class ConguateGradientBase(Transform, ABC):
 
         # resetting
         self.global_state['step'] = step + 1
-        reset_interval = self.settings[params[0]]['reset_interval']
+        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()
 
@@ -82,7 +84,7 @@ def fletcher_reeves_beta(gg: torch.Tensor, prev_gg: torch.Tensor):
 
 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 = None, clip_beta=False, inner: Chainable | None = None):
+    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):
@@ -104,7 +106,7 @@ def hestenes_stiefel_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
 
 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 = None, clip_beta=False, inner: Chainable | None = None):
+    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):
@@ -119,7 +121,7 @@ def dai_yuan_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
 
 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 = None, clip_beta=False, inner: Chainable | None = None):
+    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):
@@ -134,7 +136,7 @@ def liu_storey_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList, ):
 
 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 = None, clip_beta=False, inner: Chainable | None = None):
+    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):
@@ -151,10 +153,10 @@ class ConjugateDescent(Transform):
 
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply(self, tensors, params, grads, loss, states, settings):
         g = as_tensorlist(tensors)
 
-        prev_d = self.get_state('prev_dir', params=params, cls=TensorList, init = torch.zeros_like)
+        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])
 
@@ -164,7 +166,7 @@ class ConjugateDescent(Transform):
 
         # inner step
         if 'inner' in self.children:
-            g = as_tensorlist(apply(self.children['inner'], g, params, grads, vars))
+            g = as_tensorlist(apply_transform(self.children['inner'], g, params, grads))
 
         dir = g.add_(prev_d.mul_(beta))
         prev_d.copy_(dir)
@@ -187,7 +189,7 @@ def hager_zhang_beta(g:TensorList, prev_d:TensorList, prev_g:TensorList,):
 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 = None, clip_beta=False, inner: Chainable | None = None):
+    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):
@@ -211,8 +213,56 @@ def hs_dy_beta(g: TensorList, prev_d: TensorList,prev_g: TensorList):
 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 = None, clip_beta=False, inner: Chainable | None = None):
+    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/modular_lbfgs.py

@@ -4,7 +4,7 @@ from typing import Any
 
 import torch
 
-from ....core import Chainable, Module, Transform, Vars, apply, maybe_chain
+from ....core import Chainable, Module, Transform, Var, apply_transform, maybe_chain
 from ....utils import NumberList, TensorList, as_tensorlist
 
 
@@ -28,7 +28,7 @@ def _adaptive_damping(
 
 def lbfgs(
     tensors_: TensorList,
-    vars: Vars,
+    var: Var,
     s_history: deque[TensorList],
     y_history: deque[TensorList],
     sy_history: deque[torch.Tensor],
@@ -60,7 +60,7 @@ def lbfgs(
         z = q * (ys_k / (y_k.dot(y_k)))
 
         if z_tfm is not None:
-            z = TensorList(apply(z_tfm, tensors=z, params=vars.params, grads=vars.grad, vars=vars))
+            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)):
@@ -73,28 +73,28 @@ def lbfgs(
 def _apply_tfms_into_history(
     self: Module,
     params: list[torch.Tensor],
-    vars: Vars,
+    var: Var,
     update: list[torch.Tensor],
 ):
     if 'params_history_tfm' in self.children:
-        params = apply(self.children['params_history_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=vars.grad, vars=vars)
+        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(self.children['grad_history_tfm'], tensors=as_tensorlist(update).clone(), params=params, grads=vars.grad, vars=vars)
+        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],
-    vars: Vars,
+    var: Var,
     update: list[torch.Tensor],
 ):
     if 'params_precond_tfm' in self.children:
-        params = apply(self.children['params_precond_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=vars.grad, vars=vars)
+        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(self.children['grad_precond_tfm'], tensors=update, params=params, grads=vars.grad, vars=vars)
+        update = apply_transform(self.children['grad_precond_tfm'], tensors=update, params=params, grads=var.grad, var=var)
 
     return params, update
 
@@ -165,9 +165,9 @@ class ModularLBFGS(Module):
         self.global_state['sy_history'].clear()
 
     @torch.no_grad
-    def step(self, vars):
-        params = as_tensorlist(vars.params)
-        update = as_tensorlist(vars.get_update())
+    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
 
@@ -186,11 +186,11 @@ class ModularLBFGS(Module):
         params_h, update_h = _apply_tfms_into_history(
             self,
             params=params,
-            vars=vars,
+            var=var,
             update=update,
         )
 
-        prev_params_h, prev_grad_h = self.get_state('prev_params_h', 'prev_grad_h', params=params, cls=TensorList)
+        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:
@@ -217,16 +217,16 @@ class ModularLBFGS(Module):
         # step with inner module before applying preconditioner
         if 'update_precond_tfm' in self.children:
             update_precond_tfm = self.children['update_precond_tfm']
-            inner_vars = update_precond_tfm.step(vars.clone(clone_update=True))
-            vars.update_attrs_from_clone_(inner_vars)
-            tensors = inner_vars.update
+            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, vars=vars, update=update)
-        prev_params_p, prev_grad_p = self.get_state('prev_params_p', 'prev_grad_p', params=params, cls=TensorList)
+        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
@@ -245,13 +245,13 @@ class ModularLBFGS(Module):
         # 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:
-                vars.update = update # may have been updated by inner module, probably makes sense to use it here?
-                return vars
+                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),
-            vars=vars,
+            var=var,
             s_history=s_history,
             y_history=y_history,
             sy_history=sy_history,
@@ -260,7 +260,7 @@ class ModularLBFGS(Module):
             z_tfm=self.children.get('z_tfm', None),
         )
 
-        vars.update = dir
+        var.update = dir
 
-        return vars
+        return var
 

torchzero/modules/quasi_newton/lbfgs.py

@@ -2,7 +2,7 @@ from collections import deque
 from operator import itemgetter
 import torch
 
-from ...core import Transform, Chainable, Module, Vars, apply
+from ...core import Transform, Chainable, Module, Var, apply_transform
 from ...utils import TensorList, as_tensorlist, NumberList
 
 
@@ -154,9 +154,9 @@ class LBFGS(Module):
         self.global_state['sy_history'].clear()
 
     @torch.no_grad
-    def step(self, vars):
-        params = as_tensorlist(vars.params)
-        update = as_tensorlist(vars.get_update())
+    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
 
@@ -167,10 +167,10 @@ class LBFGS(Module):
 
         tol, damping, init_damping, eigval_bounds, update_freq, z_beta, tol_reset = itemgetter(
             'tol', 'damping', 'init_damping', 'eigval_bounds', 'update_freq', 'z_beta', 'tol_reset')(self.settings[params[0]])
-        params_beta, grads_beta = self.get_settings('params_beta', 'grads_beta', params=params)
+        params_beta, grads_beta = self.get_settings(params, 'params_beta', 'grads_beta')
 
         l_params, l_update = _lerp_params_update_(self, params, update, params_beta, grads_beta)
-        prev_l_params, prev_l_grad = self.get_state('prev_l_params', 'prev_l_grad', params=params, cls=TensorList)
+        prev_l_params, prev_l_grad = self.get_state(params, 'prev_l_params', 'prev_l_grad', cls=TensorList)
 
         # 1st step - there are no previous params and grads, `lbfgs` will do normalized SGD step
         if step == 0:
@@ -196,19 +196,19 @@ class LBFGS(Module):
 
         # step with inner module before applying preconditioner
         if self.children:
-            update = TensorList(apply(self.children['inner'], tensors=update, params=params, grads=vars.grad, vars=vars))
+            update = TensorList(apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var))
 
         # tolerance on gradient difference to avoid exploding after converging
         if tol is not None:
             if y_k is not None and y_k.abs().global_max() <= tol:
-                vars.update = update # may have been updated by inner module, probably makes sense to use it here?
+                var.update = update # may have been updated by inner module, probably makes sense to use it here?
                 if tol_reset: self.reset()
-                return vars
+                return var
 
         # lerp initial H^-1 @ q guess
         z_ema = None
         if z_beta is not None:
-            z_ema = self.get_state('z_ema', params=vars.params, cls=TensorList)
+            z_ema = self.get_state(var.params, 'z_ema', cls=TensorList)
 
         # precondition
         dir = lbfgs(
@@ -223,7 +223,7 @@ class LBFGS(Module):
             step=step
         )
 
-        vars.update = dir
+        var.update = dir
 
-        return vars
+        return var
 

torchzero/modules/quasi_newton/lsr1.py

@@ -3,7 +3,7 @@ from operator import itemgetter
 
 import torch
 
-from ...core import Chainable, Module, Transform, Vars, apply
+from ...core import Chainable, Module, Transform, Var, apply_transform
 from ...utils import NumberList, TensorList, as_tensorlist
 
 from .lbfgs import _lerp_params_update_
@@ -123,9 +123,9 @@ class LSR1(Module):
 
 
     @torch.no_grad
-    def step(self, vars: Vars):
-        params = as_tensorlist(vars.params)
-        update = as_tensorlist(vars.get_update())
+    def step(self, var: 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
 
@@ -135,10 +135,10 @@ class LSR1(Module):
         settings = self.settings[params[0]]
         tol, update_freq, scale_second = itemgetter('tol', 'update_freq', 'scale_second')(settings)
 
-        params_beta, grads_beta_ = self.get_settings('params_beta', 'grads_beta', params=params) # type: ignore
+        params_beta, grads_beta_ = self.get_settings(params, 'params_beta', 'grads_beta') # type: ignore
         l_params, l_update = _lerp_params_update_(self, params, update, params_beta, grads_beta_)
 
-        prev_l_params, prev_l_grad = self.get_state('prev_l_params', 'prev_l_grad', params=params, cls=TensorList)
+        prev_l_params, prev_l_grad = self.get_state(params, 'prev_l_params', 'prev_l_grad', cls=TensorList)
 
         y_k = None
         if step != 0:
@@ -153,13 +153,13 @@ class LSR1(Module):
         prev_l_grad.copy_(l_update)
 
         if 'inner' in self.children:
-            update = TensorList(apply(self.children['inner'], tensors=update, params=params, grads=vars.grad, vars=vars))
+            update = TensorList(apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var))
 
         # tolerance on gradient difference to avoid exploding after converging
         if tol is not None:
             if y_k is not None and y_k.abs().global_max() <= tol:
-                vars.update = update
-                return vars
+                var.update = update
+                return var
 
         dir = lsr1_(
             tensors_=update,
@@ -169,6 +169,6 @@ class LSR1(Module):
             scale_second=scale_second,
         )
 
-        vars.update = dir
+        var.update = dir
 
-        return vars
+        return var

torchzero/modules/quasi_newton/olbfgs.py

@@ -5,17 +5,17 @@ from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, Transform, Vars, apply
+from ...core import Chainable, Module, Transform, Var, apply_transform
 from ...utils import NumberList, TensorList, as_tensorlist
 from .lbfgs import _adaptive_damping, lbfgs
 
 
 @torch.no_grad
-def _store_sk_yk_after_step_hook(optimizer, vars: Vars, prev_params: TensorList, prev_grad: TensorList, damping, init_damping, eigval_bounds, s_history: deque[TensorList], y_history: deque[TensorList], sy_history: deque[torch.Tensor]):
-    assert vars.closure is not None
-    with torch.enable_grad(): vars.closure()
-    grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in vars.params]
-    s_k = vars.params - prev_params
+def _store_sk_yk_after_step_hook(optimizer, var: Var, prev_params: TensorList, prev_grad: TensorList, damping, init_damping, eigval_bounds, s_history: deque[TensorList], y_history: deque[TensorList], sy_history: deque[torch.Tensor]):
+    assert var.closure is not None
+    with torch.enable_grad(): var.closure()
+    grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in var.params]
+    s_k = var.params - prev_params
     y_k = grad - prev_grad
     ys_k = s_k.dot(y_k)
 
@@ -95,11 +95,11 @@ class OnlineLBFGS(Module):
         self.global_state['sy_history'].clear()
 
     @torch.no_grad
-    def step(self, vars):
-        assert vars.closure is not None
+    def step(self, var):
+        assert var.closure is not None
 
-        params = as_tensorlist(vars.params)
-        update = as_tensorlist(vars.get_update())
+        params = as_tensorlist(var.params)
+        update = as_tensorlist(var.get_update())
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
@@ -113,7 +113,7 @@ class OnlineLBFGS(Module):
 
         # sample gradient at previous params with current mini-batch
         if sample_grads == 'before':
-            prev_params = self.get_state('prev_params', params=params, cls=TensorList)
+            prev_params = self.get_state(params, 'prev_params', cls=TensorList)
             if step == 0:
                 s_k = None; y_k = None; ys_k = None
             else:
@@ -121,7 +121,7 @@ class OnlineLBFGS(Module):
 
                 current_params = params.clone()
                 params.set_(prev_params)
-                with torch.enable_grad(): vars.closure()
+                with torch.enable_grad(): var.closure()
                 y_k = update - params.grad
                 ys_k = s_k.dot(y_k)
                 params.set_(current_params)
@@ -146,7 +146,7 @@ class OnlineLBFGS(Module):
                 ys_k = s_k.dot(y_k)
 
             # this will run after params are updated by Modular after running all future modules
-            vars.post_step_hooks.append(
+            var.post_step_hooks.append(
                 partial(
                     _store_sk_yk_after_step_hook,
                     prev_params=params.clone(),
@@ -164,18 +164,18 @@ class OnlineLBFGS(Module):
 
         # step with inner module before applying preconditioner
         if self.children:
-            update = TensorList(apply(self.children['inner'], tensors=update, params=params, grads=vars.grad, vars=vars))
+            update = TensorList(apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var))
 
         # tolerance on gradient difference to avoid exploding after converging
         if tol is not None:
             if y_k is not None and y_k.abs().global_max() <= tol:
-                vars.update = update # may have been updated by inner module, probably makes sense to use it here?
-                return vars
+                var.update = update # may have been updated by inner module, probably makes sense to use it here?
+                return var
 
         # lerp initial H^-1 @ q guess
         z_ema = None
         if z_beta is not None:
-            z_ema = self.get_state('z_ema', params=vars.params, cls=TensorList)
+            z_ema = self.get_state(params, 'z_ema', cls=TensorList)
 
         # precondition
         dir = lbfgs(
@@ -190,7 +190,7 @@ class OnlineLBFGS(Module):
             step=step
         )
 
-        vars.update = dir
+        var.update = dir
 
-        return vars
+        return var