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

tests/test_vars.py

@@ -1,10 +1,10 @@
 import pytest
 import torch
-from torchzero.core.module import Vars
+from torchzero.core.module import Var
 from torchzero.utils.tensorlist import TensorList
 
 @torch.no_grad
-def test_vars_get_loss():
+def test_var_get_loss():
 
     # ---------------------------- test that it works ---------------------------- #
     params = [torch.tensor(2.0, requires_grad=True)]
@@ -26,20 +26,20 @@ def test_vars_get_loss():
             assert not loss.requires_grad, "loss requires grad with backward=False"
         return loss
 
-    vars = Vars(params=params, closure=closure_1, model=None, current_step=0)
+    var = Var(params=params, closure=closure_1, model=None, current_step=0)
 
-    assert vars.loss is None, vars.loss
+    assert var.loss is None, var.loss
 
-    assert (loss := vars.get_loss(backward=False)) == 4.0, loss
+    assert (loss := var.get_loss(backward=False)) == 4.0, loss
     assert evaluated, evaluated
-    assert loss is vars.loss
-    assert vars.loss == 4.0
-    assert vars.loss_approx == 4.0
-    assert vars.grad is None, vars.grad
+    assert loss is var.loss
+    assert var.loss == 4.0
+    assert var.loss_approx == 4.0
+    assert var.grad is None, var.grad
 
     # reevaluate, which should just return already evaluated loss
-    assert (loss := vars.get_loss(backward=False)) == 4.0, loss
-    assert vars.grad is None, vars.grad
+    assert (loss := var.get_loss(backward=False)) == 4.0, loss
+    assert var.grad is None, var.grad
 
 
     # ----------------------- test that backward=True works ---------------------- #
@@ -61,30 +61,30 @@ def test_vars_get_loss():
             assert not loss.requires_grad, "loss requires grad with backward=False"
         return loss
 
-    vars = Vars(params=params, closure=closure_2, model=None, current_step=0)
-    assert vars.grad is None, vars.grad
-    assert (loss := vars.get_loss(backward=True)) == 6.0, loss
-    assert vars.grad is not None
-    assert vars.grad[0] == 2.0, vars.grad
+    var = Var(params=params, closure=closure_2, model=None, current_step=0)
+    assert var.grad is None, var.grad
+    assert (loss := var.get_loss(backward=True)) == 6.0, loss
+    assert var.grad is not None
+    assert var.grad[0] == 2.0, var.grad
 
     # reevaluate, which should just return already evaluated loss
-    assert (loss := vars.get_loss(backward=True)) == 6.0, loss
-    assert vars.grad[0] == 2.0, vars.grad
+    assert (loss := var.get_loss(backward=True)) == 6.0, loss
+    assert var.grad[0] == 2.0, var.grad
 
     # get grad, which should just return already evaluated grad
-    assert (grad := vars.get_grad())[0] == 2.0, grad
-    assert grad is vars.grad, grad
+    assert (grad := var.get_grad())[0] == 2.0, grad
+    assert grad is var.grad, grad
 
     # get update, which should create and return cloned grad
-    assert vars.update is None
-    assert (update := vars.get_update())[0] == 2.0, update
-    assert update is vars.update
-    assert update is not vars.grad
-    assert vars.grad is not None
-    assert update[0] == vars.grad[0]
+    assert var.update is None
+    assert (update := var.get_update())[0] == 2.0, update
+    assert update is var.update
+    assert update is not var.grad
+    assert var.grad is not None
+    assert update[0] == var.grad[0]
 
 @torch.no_grad
-def test_vars_get_grad():
+def test_var_get_grad():
     params = [torch.tensor(2.0, requires_grad=True)]
     evaluated = False
 
@@ -103,20 +103,20 @@ def test_vars_get_grad():
             assert not loss.requires_grad, "loss requires grad with backward=False"
         return loss
 
-    vars = Vars(params=params, closure=closure, model=None, current_step=0)
-    assert (grad := vars.get_grad())[0] == 4.0, grad
-    assert grad is vars.grad
+    var = Var(params=params, closure=closure, model=None, current_step=0)
+    assert (grad := var.get_grad())[0] == 4.0, grad
+    assert grad is var.grad
 
-    assert vars.loss == 4.0
-    assert (loss := vars.get_loss(backward=False)) == 4.0, loss
-    assert (loss := vars.get_loss(backward=True)) == 4.0, loss
-    assert vars.loss_approx == 4.0
+    assert var.loss == 4.0
+    assert (loss := var.get_loss(backward=False)) == 4.0, loss
+    assert (loss := var.get_loss(backward=True)) == 4.0, loss
+    assert var.loss_approx == 4.0
 
-    assert vars.update is None, vars.update
-    assert (update := vars.get_update())[0] == 4.0, update
+    assert var.update is None, var.update
+    assert (update := var.get_update())[0] == 4.0, update
 
 @torch.no_grad
-def test_vars_get_update():
+def test_var_get_update():
     params = [torch.tensor(2.0, requires_grad=True)]
     evaluated = False
 
@@ -135,27 +135,28 @@ def test_vars_get_update():
             assert not loss.requires_grad, "loss requires grad with backward=False"
         return loss
 
-    vars = Vars(params=params, closure=closure, model=None, current_step=0)
-    assert vars.update is None, vars.update
-    assert (update := vars.get_update())[0] == 4.0, update
-    assert update is vars.update
+    var = Var(params=params, closure=closure, model=None, current_step=0)
+    assert var.update is None, var.update
+    assert (update := var.get_update())[0] == 4.0, update
+    assert update is var.update
 
-    assert (grad := vars.get_grad())[0] == 4.0, grad
-    assert grad is vars.grad
+    assert (grad := var.get_grad())[0] == 4.0, grad
+    assert grad is var.grad
     assert grad is not update
 
-    assert vars.loss == 4.0
-    assert (loss := vars.get_loss(backward=False)) == 4.0, loss
-    assert (loss := vars.get_loss(backward=True)) == 4.0, loss
-    assert vars.loss_approx == 4.0
+    assert var.loss == 4.0
+    assert (loss := var.get_loss(backward=False)) == 4.0, loss
+    assert (loss := var.get_loss(backward=True)) == 4.0, loss
+    assert var.loss_approx == 4.0
 
-    assert (update := vars.get_update())[0] == 4.0, update
+    assert (update := var.get_update())[0] == 4.0, update
 
 
-def _assert_vars_are_same_(v1: Vars, v2: Vars, clone_update: bool):
+def _assert_var_are_same_(v1: Var, v2: Var, clone_update: bool):
     for k,v in v1.__dict__.items():
         if not k.startswith('__'):
             # if k == 'post_step_hooks': continue
+            if k == 'storage': continue
             if k == 'update' and clone_update:
                 if v1.update is None or v2.update is None:
                     assert v1.update is None and v2.update is None, f'{k} is not the same, {v1 = }, {v2 = }'
@@ -165,20 +166,20 @@ def _assert_vars_are_same_(v1: Vars, v2: Vars, clone_update: bool):
             else:
                 assert getattr(v2, k) is v, f'{k} is not the same, {v1 = }, {v2 = }'
 
-def test_vars_clone():
+def test_var_clone():
     model = torch.nn.Sequential(torch.nn.Linear(2,2), torch.nn.Linear(2,4))
     def closure(backward): return 1
-    vars = Vars(params=list(model.parameters()), closure=closure, model=model, current_step=0)
+    var = Var(params=list(model.parameters()), closure=closure, model=model, current_step=0)
 
-    _assert_vars_are_same_(vars, vars.clone(clone_update=False), clone_update=False)
-    _assert_vars_are_same_(vars, vars.clone(clone_update=True), clone_update=True)
+    _assert_var_are_same_(var, var.clone(clone_update=False), clone_update=False)
+    _assert_var_are_same_(var, var.clone(clone_update=True), clone_update=True)
 
-    vars.grad = TensorList(torch.randn(5))
-    _assert_vars_are_same_(vars, vars.clone(clone_update=False), clone_update=False)
-    _assert_vars_are_same_(vars, vars.clone(clone_update=True), clone_update=True)
+    var.grad = TensorList(torch.randn(5))
+    _assert_var_are_same_(var, var.clone(clone_update=False), clone_update=False)
+    _assert_var_are_same_(var, var.clone(clone_update=True), clone_update=True)
 
-    vars.update = TensorList(torch.randn(5) * 2)
-    vars.loss = torch.randn(1)
-    vars.loss_approx = vars.loss
-    _assert_vars_are_same_(vars, vars.clone(clone_update=False), clone_update=False)
-    _assert_vars_are_same_(vars, vars.clone(clone_update=True), clone_update=True)
+    var.update = TensorList(torch.randn(5) * 2)
+    var.loss = torch.randn(1)
+    var.loss_approx = var.loss
+    _assert_var_are_same_(var, var.clone(clone_update=False), clone_update=False)
+    _assert_var_are_same_(var, var.clone(clone_update=True), clone_update=True)

torchzero/core/__init__.py

@@ -1,3 +1,2 @@
-from .module import Vars, Module, Modular, Chain, maybe_chain, Chainable
-from .transform import Transform, TensorwiseTransform, Target, apply
-from .preconditioner import Preconditioner, TensorwisePreconditioner
+from .module import Var, Module, Modular, Chain, maybe_chain, Chainable
+from .transform import Transform, TensorwiseTransform, Target, apply_transform

torchzero/core/module.py

@@ -3,7 +3,7 @@ from abc import ABC, abstractmethod
 from collections import ChainMap, defaultdict
 from collections.abc import Callable, Iterable, MutableMapping, Sequence
 from operator import itemgetter
-from typing import Any, final, overload
+from typing import Any, final, overload, Literal
 
 import torch
 
@@ -14,6 +14,7 @@ from ..utils import (
     _make_param_groups,
     get_state_vals,
 )
+from ..utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
 from ..utils.python_tools import flatten
 
 
@@ -29,8 +30,8 @@ def _closure_backward(closure, params, retain_graph, create_graph):
         return loss
 
 # region Vars
-# ----------------------------------- vars ----------------------------------- #
-class Vars:
+# ----------------------------------- var ----------------------------------- #
+class Var:
     """
     Holds the state and context passed between optimizer modules during a step.
 
@@ -74,13 +75,13 @@ class Vars:
         """loss at a point near current point. This can be useful as some modules only calculate loss at perturbed points,
         whereas some other modules require loss strictly at current point."""
 
-        self.post_step_hooks: list[Callable[[Modular, Vars]]] = []
+        self.post_step_hooks: list[Callable[[Modular, Var]]] = []
         """list of functions to be called after optimizer step.
         The signature is:
 
         .. code:: py
 
-            def hook(optimizer: Modular, vars: Vars): ...
+            def hook(optimizer: Modular, var: Vars): ...
 
         """
 
@@ -109,8 +110,11 @@ class Vars:
         self.skip_update: bool = False
         """if True, the parameters will not be updated"""
 
+        self.storage: dict = {}
+        """Storage for any other data, such as hessian estimates, etc"""
+
     def get_loss(self, backward: bool, retain_graph = None, create_graph: bool = False) -> torch.Tensor | float:
-        """Returns the loss at current parameters, computing it if it hasn't been computed already and assigning :code:`vars.loss`.
+        """Returns the loss at current parameters, computing it if it hasn't been computed already and assigning :code:`var.loss`.
         Do not call this at perturbed parameters. Backward always zeroes grads before recomputing."""
 
         if self.loss is None:
@@ -143,7 +147,7 @@ class Vars:
 
     def get_grad(self, retain_graph: bool | None = None, create_graph: bool = False) -> list[torch.Tensor]:
         """Returns the gradient at initial parameters, computing it if it hasn't been computed already and assigning
-        :code:`vars.grad` and potentially :code:`vars.loss`. Do not call this at perturbed parameters."""
+        :code:`var.grad` and potentially :code:`var.loss`. Do not call this at perturbed parameters."""
         if self.grad is None:
             if self.closure is None: raise RuntimeError("closure is None")
             self.get_loss(backward=True, retain_graph=retain_graph, create_graph=create_graph) # evaluate and set self.loss and self.grad
@@ -152,15 +156,15 @@ class Vars:
         return self.grad
 
     def get_update(self) -> list[torch.Tensor]:
-        """Returns the update. If update is None, it is initialized by cloning the gradients and assigning to :code:`vars.update`.
-        Computing the gradients may assign :code:`vars.grad` and :code:`vars.loss` if they haven't been computed.
+        """Returns the update. If update is None, it is initialized by cloning the gradients and assigning to :code:`var.update`.
+        Computing the gradients may assign :code:`var.grad` and :code:`var.loss` if they haven't been computed.
         Do not call this at perturbed parameters."""
         if self.update is None: self.update = [g.clone() for g in self.get_grad()]
         return self.update
 
     def clone(self, clone_update: bool):
         """Creates a shallow copy of the Vars object, update can optionally be deep-copied (via :code:`torch.clone`)."""
-        copy = Vars(params = self.params, closure=self.closure, model=self.model, current_step=self.current_step)
+        copy = Var(params = self.params, closure=self.closure, model=self.model, current_step=self.current_step)
 
         if clone_update and self.update is not None:
             copy.update = [u.clone() for u in self.update]
@@ -176,16 +180,17 @@ class Vars:
 
         return copy
 
-    def update_attrs_from_clone_(self, vars: "Vars"):
+    def update_attrs_from_clone_(self, var: "Var"):
         """Updates attributes of this `Vars` instance from a cloned instance.
         Typically called after a child module has processed a cloned `Vars`
         object. This propagates any newly computed loss or gradient values
         from the child's context back to the parent `Vars` if the parent
         didn't have them computed already.
         """
-        if self.loss is None: self.loss = vars.loss
-        if self.loss_approx is None: self.loss_approx = vars.loss_approx
-        if self.grad is None: self.grad = vars.grad
+        if self.loss is None: self.loss = var.loss
+        if self.loss_approx is None: self.loss_approx = var.loss_approx
+        if self.grad is None: self.grad = var.grad
+        self.storage.update(var.storage)
 
     def zero_grad(self, set_to_none=True):
         if set_to_none:
@@ -269,36 +274,36 @@ class Module(ABC):
         return s
 
     @overload
-    def get_settings(self, key: str, *,
-                     params: Sequence[torch.Tensor], cls: type[ListLike] = list) -> ListLike: ...
+    def get_settings(self, params: Sequence[torch.Tensor], key: str, *,
+                     cls: type[ListLike] = list) -> ListLike: ...
     @overload
-    def get_settings(self, key: list[str] | tuple[str,...], *,
-                     params: Sequence[torch.Tensor], cls: type[ListLike] = list) -> list[ListLike]: ...
+    def get_settings(self, params: Sequence[torch.Tensor], key: list[str] | tuple[str,...], *,
+                     cls: type[ListLike] = list) -> list[ListLike]: ...
     @overload
-    def get_settings(self, key: str, key2: str, *keys: str,
-                     params: Sequence[torch.Tensor], cls: type[ListLike] = list) -> list[ListLike]: ...
+    def get_settings(self, params: Sequence[torch.Tensor], key: str, key2: str, *keys: str,
+                     cls: type[ListLike] = list) -> list[ListLike]: ...
 
-    def get_settings(self, key: str | list[str] | tuple[str,...], key2: str | None = None, *keys: str,
-                     params: Sequence[torch.Tensor], cls: type[ListLike] = list) -> ListLike | list[ListLike]:
+    def get_settings(self, params: Sequence[torch.Tensor], key: str | list[str] | tuple[str,...], key2: str | None = None,
+                     *keys: str, cls: type[ListLike] = list) -> ListLike | list[ListLike]:
         # if isinstance(params, Vars): params = params.params
         return get_state_vals(self.settings, params, key, key2, *keys, must_exist=True, cls=cls) # pyright:ignore[reportArgumentType]
 
 
     @overload
-    def get_state(self, key: str, *,
-                   params: Sequence[torch.Tensor], must_exist: bool = False, init: Init = torch.zeros_like,
+    def get_state(self, params: Sequence[torch.Tensor], key: str, *,
+                   must_exist: bool = False, init: Init = torch.zeros_like,
                    cls: type[ListLike] = list) -> ListLike: ...
     @overload
-    def get_state(self, key: list[str] | tuple[str,...], *,
-                   params: Sequence[torch.Tensor], must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+    def get_state(self, params: Sequence[torch.Tensor], key: list[str] | tuple[str,...], *,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
                    cls: type[ListLike] = list) -> list[ListLike]: ...
     @overload
-    def get_state(self, key: str, key2: str, *keys: str,
-                   params: Sequence[torch.Tensor], must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+    def get_state(self, params: Sequence[torch.Tensor], key: str, key2: str, *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
                    cls: type[ListLike] = list) -> list[ListLike]: ...
 
-    def get_state(self, key: str | list[str] | tuple[str,...], key2: str | None = None, *keys: str,
-                   params: Sequence[torch.Tensor], must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+    def get_state(self, params: Sequence[torch.Tensor], key: str | list[str] | tuple[str,...], key2: str | None = None, *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
                    cls: type[ListLike] = list) -> ListLike | list[ListLike]:
         """Returns values of per-parameter state for a given key.
         If key doesn't exist, create it with inits.
@@ -358,6 +363,26 @@ class Module(ABC):
     #     # if isinstance(params, Vars): params = params.params
     #     return itemgetter(*keys)(self.settings[params[0]])
 
+    def clear_state_keys(self, *keys:str):
+        for s in self.state.values():
+            for k in keys:
+                if k in s: del s[k]
+
+    @overload
+    def store(self, params: Sequence[torch.Tensor], keys: str, values: Sequence): ...
+    @overload
+    def store(self, params: Sequence[torch.Tensor], keys: Sequence[str], values: Sequence[Sequence]): ...
+    def store(self, params: Sequence[torch.Tensor], keys: str | Sequence[str], values: Sequence):
+        if isinstance(keys, str):
+            for p,v in zip(params, values):
+                state = self.state[p]
+                state[keys] = v
+            return
+
+        for p, *p_v in zip(params, *values):
+            state = self.state[p]
+            for k,v in zip(keys, p_v): state[k] = v
+
     def state_dict(self):
         """state dict"""
         packed_state = {id(k):v for k,v in self.state.items()}
@@ -403,23 +428,111 @@ class Module(ABC):
         self._extra_unpack(state_dict['extra'])
 
     # ---------------------------- OVERRIDABLE METHODS --------------------------- #
-    @abstractmethod
-    def step(self, vars: Vars) -> Vars:
-        """performs a step, returns new vars but may update them in-place."""
+    def step(self, var: Var) -> Var:
+        """performs a step, returns new var but may update it in-place."""
+        self.update(var)
+        return self.apply(var)
+
+    def update(self, var:Var) -> Any:
+        """Updates the internal state of this module. This should not modify `var.update`.
+
+        Specifying ``update`` and ``apply`` methods is optional and allows certain meta-modules to be used,
+        such as ::code::`tz.m.Online`.
+        """
+
+    def apply(self, var: Var) -> Var:
+        """Applies this module to ``var.get_update()``. This should not modify the internal state of this module if possible."""
+        raise NotImplementedError(f"{self} doesn't implement the `apply` method.")
 
     def reset(self):
-        """Resets the internal state of the module (e.g. momentum)."""
+        """Resets the internal state of the module (e.g. momentum). By default clears state and global state."""
         # no complex logic is allowed there because this is overridden by many modules
         # where super().reset() shouldn't be called
         self.state.clear()
         self.global_state.clear()
 
+    def reset_for_online(self):
+        """resets only the intermediate state of this module, e.g. previous parameters and gradient."""
+        for c in self.children.values(): c.reset_for_online()
+
     def _extra_pack(self):
         return {}
 
     def _extra_unpack(self, x):
         pass
 
+
+    # ------------------------------ HELPER METHODS ------------------------------ #
+    @torch.no_grad
+    def Hvp(
+        self,
+        v: Sequence[torch.Tensor],
+        at_x0: bool,
+        var: Var,
+        rgrad: Sequence[torch.Tensor] | None,
+        hvp_method: Literal['autograd', 'forward', 'central'],
+        h: float,
+        normalize: bool,
+        retain_grad: bool,
+    ):
+        """
+        Returns ``(Hvp, rgrad)``. ``rgrad`` is gradient at current parameters, possibly with create_graph=True, or it may be None with ``hvp_method="central"``. Gradient is set to vars automatically if ``at_x0``, you can always access it with ``vars.get_grad()``
+
+        Single sample example:
+
+        .. code:: py
+
+            Hvp, _ = self.hvp(v, at_x0=True, rgrad=None, ..., retain_graph=False)
+
+        Multiple samples example:
+
+        .. code:: py
+
+            D = None
+            rgrad = None
+            for i in range(n_samples):
+                v = [torch.randn_like(p) for p in params]
+                Hvp, rgrad = self.hvp(v, at_x0=True, rgrad=rgrad, ..., retain_graph=i < n_samples-1)
+
+                if D is None: D = Hvp
+                else: torch._foreach_add_(D, Hvp)
+
+            if n_samples > 1: torch._foreach_div_(D, n_samples)
+        Args:
+            v (Sequence[torch.Tensor]): vector in hessian-vector product
+            at_x0 (bool): whether this is being called at original or perturbed parameters.
+            var (Var): Var
+            rgrad (Sequence[torch.Tensor] | None): pass None initially, then pass what this returns.
+            hvp_method (str): hvp method.
+            h (float): finite difference step size
+            normalize (bool): whether to normalize v for finite difference
+            retain_grad (bool): retain grad
+        """
+        # get grad
+        if rgrad is None and hvp_method in ('autograd', 'forward'):
+            if at_x0: rgrad = var.get_grad(create_graph = hvp_method=='autograd')
+            else:
+                if var.closure is None: raise RuntimeError("Closure is required to calculate HVp")
+                with torch.enable_grad():
+                    loss = var.closure()
+                    rgrad = torch.autograd.grad(loss, var.params, create_graph = hvp_method=='autograd')
+
+        if hvp_method == 'autograd':
+            assert rgrad is not None
+            Hvp = hvp(var.params, rgrad, v, retain_graph=retain_grad)
+
+        elif hvp_method == 'forward':
+            assert rgrad is not None
+            loss, Hvp = hvp_fd_forward(var.closure, var.params, v, h=h, g_0=rgrad, normalize=normalize)
+
+        elif hvp_method == 'central':
+            loss, Hvp = hvp_fd_central(var.closure, var.params, v, h=h, normalize=normalize)
+
+        else:
+            raise ValueError(hvp_method)
+
+        return Hvp, rgrad
+
 # endregion
 
 Chainable = Module | Sequence[Module]
@@ -440,6 +553,21 @@ def unroll_modules(*modules: Chainable) -> list[Module]:
 
 # region Modular
 # ---------------------------------- Modular --------------------------------- #
+
+class _EvalCounterClosure:
+    """keeps track of how many times closure has been evaluated"""
+    __slots__ = ("modular", "closure")
+    def __init__(self, modular: "Modular", closure):
+        self.modular = modular
+        self.closure = closure
+
+    def __call__(self, *args, **kwargs):
+        if self.closure is None:
+            raise RuntimeError("One of the modules requires closure to be passed to the step method")
+
+        self.modular.num_evaluations += 1
+        return self.closure(*args, **kwargs)
+
 # have to inherit from Modular to support lr schedulers
 # although Accelerate doesn't work due to converting param_groups to a dict
 class Modular(torch.optim.Optimizer):
@@ -496,7 +624,10 @@ class Modular(torch.optim.Optimizer):
         #     self.add_param_group(param_group)
 
         self.current_step = 0
-        """The global step counter for the optimizer."""
+        """global step counter for the optimizer."""
+
+        self.num_evaluations = 0
+        """number of times the objective has been evaluated (number of closure calls or number of steps if closure is None)."""
 
     def add_param_group(self, param_group: dict[str, Any]):
         proc_param_group = _make_param_groups([param_group], differentiable=False)[0]
@@ -556,13 +687,14 @@ class Modular(torch.optim.Optimizer):
                 if not p.requires_grad: continue
                 for map in self._per_parameter_global_settings[p]: map.update(settings)
 
-        # create vars
+        # create var
         params = [p for g in self.param_groups for p in g['params'] if p.requires_grad]
-        vars = Vars(params=params, closure=closure, model=self.model, current_step=self.current_step)
+        var = Var(params=params, closure=_EvalCounterClosure(self, closure), model=self.model, current_step=self.current_step)
 
         # if closure is None, assume backward has been called and gather grads
         if closure is None:
-            vars.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+            var.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+            self.num_evaluations += 1
 
         last_module = self.modules[-1]
         last_lr = last_module.defaults.get('lr', None)
@@ -570,27 +702,27 @@ class Modular(torch.optim.Optimizer):
 
         # step
         for i, module in enumerate(self.modules):
-            if i!=0: vars = vars.clone(clone_update=False)
+            if i!=0: var = var.clone(clone_update=False)
 
             # last module, or next to last module before lr
             if (i == n_modules - 1) or ((i == n_modules - 2) and (last_lr is not None)):
-                if module.children: vars.nested_is_last = True
-                else: vars.is_last = True
-                if last_lr is not None: vars.last_module_lrs = last_module.get_settings('lr', params=vars.params)
+                if module.children: var.nested_is_last = True
+                else: var.is_last = True
+                if last_lr is not None: var.last_module_lrs = [last_module.settings[p]['lr'] for p in var.params]
 
-            vars = module.step(vars)
-            if vars.stop: break
+            var = module.step(var)
+            if var.stop: break
 
         # apply update
-        if not vars.skip_update:
+        if not var.skip_update:
             with torch.no_grad():
-                torch._foreach_sub_(params, vars.get_update())
+                torch._foreach_sub_(params, var.get_update())
 
-        for hook in vars.post_step_hooks:
-            hook(self, vars)
+        for hook in var.post_step_hooks:
+            hook(self, var)
 
         self.current_step += 1
-        return vars.loss if vars.loss is not None else vars.loss_approx
+        return var.loss if var.loss is not None else var.loss_approx
 
     def __repr__(self):
         return f'Modular({", ".join(str(m) for m in self.modules)})'
@@ -606,11 +738,11 @@ class Chain(Module):
         for i, module in enumerate(flat_modules):
             self.set_child(f'module_{i}', module)
 
-    def step(self, vars):
+    def step(self, var):
         for i in range(len(self.children)):
-            vars = self.children[f'module_{i}'].step(vars)
-            if vars.stop: break
-        return vars
+            var = self.children[f'module_{i}'].step(var)
+            if var.stop: break
+        return var
 
     def __repr__(self):
         s = self.__class__.__name__