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

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, cast
 
 import torch
 
@@ -14,7 +14,9 @@ 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
+from ..utils.linalg.linear_operator import LinearOperator
 
 
 def _closure_backward(closure, params, retain_graph, create_graph):
@@ -32,11 +34,9 @@ def _closure_backward(closure, params, retain_graph, create_graph):
 # ----------------------------------- var ----------------------------------- #
 class Var:
     """
-    Holds the state and context passed between optimizer modules during a step.
+    Holds parameters, gradient, update, objective function (closure) if supplied, loss, and some other info.
+    Modules take in a ``Var`` object, modify and it is passed to the next module.
 
-    This class acts as a mutable container for information relevant to the current
-    optimization step, such as parameters, gradients, loss, and the computed update.
-    Modules read from and write to this object to coordinate their actions.
     """
     def __init__(
         self,
@@ -44,6 +44,10 @@ class Var:
         closure: Callable | None,
         model: torch.nn.Module | None,
         current_step: int,
+        parent: "Var | None" = None,
+        modular: "Modular | None" = None,
+        loss: torch.Tensor | None = None,
+        storage: dict | None = None,
     ):
         self.params: list[torch.Tensor] = params
         """List of all parameters with requires_grad = True."""
@@ -55,19 +59,31 @@ class Var:
         """torch.nn.Module object of the model, None if it wasn't specified."""
 
         self.current_step: int = current_step
-        """global current step, starts at 0"""
+        """global current step, starts at 0. This may not correspond to module current step,
+        for example a module may step every 10 global steps."""
+
+        self.parent: "Var | None" = parent
+        """parent ``Var`` object. When ``self.get_grad()`` is called, it will also set ``parent.grad``.
+        Same with ``self.get_loss()``. This is useful when ``self.params`` are different from ``parent.params``,
+        e.g. when projecting."""
+
+        self.modular: "Modular" = cast(Modular, modular)
+        """Modular optimizer object that created this ``Var``."""
 
         self.update: list[torch.Tensor] | None = None
         """
-        current update, at the end this is subtracted from model parameters unless it is None.
+        current update. Update is assumed to be a transformed gradient, therefore it is subtracted.
 
         If closure is None, this is initially set to cloned gradient. Otherwise this is set to None.
+
+        At the end ``var.get_update()`` is subtracted from parameters. Therefore if ``var.update`` is ``None``,
+        gradient will be used and calculated if needed.
         """
 
         self.grad: list[torch.Tensor] | None = None
-        """gradient with current parameters. If closure is not None, this is set to None and can be calculated if needed."""
+        """gradient with current parameters. If closure is not ``None``, this is set to ``None`` and can be calculated if needed."""
 
-        self.loss: torch.Tensor | Any | None = None
+        self.loss: torch.Tensor | Any | None = loss
         """loss with current parameters."""
 
         self.loss_approx: torch.Tensor | Any | None = None
@@ -76,24 +92,28 @@ class Var:
 
         self.post_step_hooks: list[Callable[[Modular, Var]]] = []
         """list of functions to be called after optimizer step.
-        The signature is:
 
-        .. code:: py
+        This attribute should always be modified in-place (using ``append`` or ``extend``).
 
-            def hook(optimizer: Modular, var: Vars): ...
+        The signature is:
 
+        ```python
+        def hook(optimizer: Modular, var: Vars): ...
+        ```
         """
 
         self.is_last: bool = False
         """
         Indicates that current module is either last or next-to-last before a learning rate module.
         This is always False if current module has children or is a child.
+        This is because otherwise the ``is_last`` would be passed to child modules, even though they aren't last.
         """
 
         self.nested_is_last: bool = False
         """
         Indicates that current module is either last or next-to-last before a learning rate module, for modules
-        that have children.
+        that have children. This will be passed to the children unless ``var.clone()`` is used, therefore
+        a child of a last module may also receive ``var.nested_is_last=True``.
         """
 
         self.last_module_lrs: list[float] | None = None
@@ -104,16 +124,30 @@ class Var:
         """
 
         self.stop: bool = False
-        """if True, all following modules will be skipped."""
+        """if True, all following modules will be skipped.
+        If this module is a child, it only affects modules at the same level (in the same Chain)."""
 
         self.skip_update: bool = False
-        """if True, the parameters will not be updated"""
+        """if True, the parameters will not be updated."""
 
-    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:`var.loss`.
-        Do not call this at perturbed parameters. Backward always zeroes grads before recomputing."""
+        # self.storage: dict = {}
+        # """Storage for any other data, such as hessian estimates, etc."""
 
+        self.attrs: dict = {}
+        """attributes, Modular.attrs is updated with this after each step. This attribute should always be modified in-place"""
+
+        if storage is None: storage = {}
+        self.storage: dict = storage
+        """additional kwargs passed to closure will end up in this dict. This attribute should always be modified in-place"""
+
+        self.should_terminate: bool | None = None
+        """termination criteria, Modular.should_terminate is set to this after each step if not None"""
+
+    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 ``var.loss``.
+        Do not call this at perturbed parameters. Backward always sets grads to None before recomputing."""
         if self.loss is None:
+
             if self.closure is None: raise RuntimeError("closure is None")
             if backward:
                 with torch.enable_grad():
@@ -124,7 +158,10 @@ class Var:
                 # initializing to zeros_like is equivalent to using zero_grad with set_to_none = False.
                 # it is technically a more correct approach for when some parameters conditionally receive gradients
                 # and in this case it shouldn't be slower.
-                self.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in self.params]
+
+                # next time closure() is called, it will set grad to None.
+                # zero_grad(set_to_none=False) shouldn't be used (I should add a warning)
+                self.grad = [p.grad if p.grad  is not None else torch.zeros_like(p) for p in self.params]
             else:
                 self.loss = self.loss_approx = self.closure(False)
 
@@ -139,11 +176,24 @@ class Var:
                     closure=self.closure, params=self.params, retain_graph=retain_graph, create_graph=create_graph
                 )
             self.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in self.params]
+
+        # set parent grad
+        if self.parent is not None:
+            # the way projections/split work, they make a new closure which evaluates original
+            # closure and projects the gradient, and set it as their var.closure.
+            # then on `get_loss(backward=True)` it is called, so it also sets original parameters gradient.
+            # and we set it to parent var here.
+            if self.parent.loss is None: self.parent.loss = self.loss
+            if self.parent.grad is None and backward:
+                if all(p.grad is None for p in self.parent.params):
+                    warnings.warn("Parent grad is None after backward.")
+                self.parent.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in self.parent.params]
+
         return self.loss # type:ignore
 
     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:`var.grad` and potentially :code:`var.loss`. Do not call this at perturbed parameters."""
+        ``var.grad`` and potentially ``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 +202,21 @@ class Var:
         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:`var.update`.
-        Computing the gradients may assign :code:`var.grad` and :code:`var.loss` if they haven't been computed.
+        """Returns the update. If update is None, it is initialized by cloning the gradients and assigning to ``var.update``.
+        Computing the gradients may assign ``var.grad`` and ``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 = Var(params = self.params, closure=self.closure, model=self.model, current_step=self.current_step)
+    def clone(self, clone_update: bool, parent: "Var | None" = None):
+        """Creates a shallow copy of the Vars object, update can optionally be deep-copied (via ``torch.clone``).
+
+        Doesn't copy ``is_last``, ``nested_is_last`` and ``last_module_lrs``. They will always be ``False``/``None``.
+
+        Setting ``parent`` is only if clone's parameters are something different,
+        while clone's closure referes to the same objective but with a "view" on parameters.
+        """
+        copy = Var(params = self.params, closure=self.closure, model=self.model, current_step=self.current_step, parent=parent)
 
         if clone_update and self.update is not None:
             copy.update = [u.clone() for u in self.update]
@@ -170,10 +226,16 @@ class Var:
         copy.grad = self.grad
         copy.loss = self.loss
         copy.loss_approx = self.loss_approx
+        copy.closure = self.closure
         copy.post_step_hooks = self.post_step_hooks
         copy.stop = self.stop
         copy.skip_update = self.skip_update
 
+        copy.modular = self.modular
+        copy.attrs = self.attrs
+        copy.storage = self.storage
+        copy.should_terminate = self.should_terminate
+
         return copy
 
     def update_attrs_from_clone_(self, var: "Var"):
@@ -182,11 +244,16 @@ class Var:
         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.
+
+        Also, as long as ``post_step_hooks`` and ``attrs`` are modified in-place,
+        if the child updates them, the update will affect the parent too.
         """
         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
 
+        if var.should_terminate is not None: self.should_terminate = var.should_terminate
+
     def zero_grad(self, set_to_none=True):
         if set_to_none:
             for p in self.params: p.grad = None
@@ -196,6 +263,7 @@ class Var:
 
 # endregion
 
+
 # region Module
 # ---------------------------------- module ---------------------------------- #
 class Module(ABC):
@@ -308,17 +376,16 @@ class Module(ABC):
 
         If you want to force it to return a tuple even with a single key, pass a list/tuple of 1 or more keys.
 
-        .. code:: py
-
-            exp_avg = self.state_vals("exp_avg")
-            # returns cls (by default TensorList)
-
-            exp_avg, exp_avg_sq = self.state_vals("exp_avg", "exp_avg_sq")
-            # returns list of cls
+        ```python
+        exp_avg = self.state_vals("exp_avg")
+        # returns cls (by default TensorList)
 
-            exp_avg = self.state_vals(["exp_avg"])
-            # always returns a list of cls, even if got a single key
+        exp_avg, exp_avg_sq = self.state_vals("exp_avg", "exp_avg_sq")
+        # returns list of cls
 
+        exp_avg = self.state_vals(["exp_avg"])
+        # always returns a list of cls, even if got a single key
+        ```
 
         Args:
             *keys (str):
@@ -358,6 +425,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()}
@@ -377,7 +464,8 @@ class Module(ABC):
         }
         return state_dict
 
-    def load_state_dict(self, state_dict: dict[str, Any], id_to_tensor: dict[int, torch.Tensor]):
+    def _load_state_dict(self, state_dict: dict[str, Any], id_to_tensor: dict[int, torch.Tensor]):
+        """loads state_dict, ``id_to_tensor`` is passed by ``Modular``"""
         # load state
         state = state_dict['state']
         self.state.clear()
@@ -396,29 +484,159 @@ class Module(ABC):
 
         # children
         for k, v in state_dict['children']:
-            if k in self.children: self.children[k].load_state_dict(v, id_to_tensor)
+            if k in self.children: self.children[k]._load_state_dict(v, id_to_tensor)
             else: warnings.warn(f'State dict for {self} has child {k}, which is missing in {self}')
 
         # extra info
         self._extra_unpack(state_dict['extra'])
 
     # ---------------------------- OVERRIDABLE METHODS --------------------------- #
-    @abstractmethod
     def step(self, var: Var) -> Var:
-        """performs a step, returns new var but may update them in-place."""
+        """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 ``tz.m.Online`` or trust regions. Alternatively, simply override the ``step`` method.
+        """
+
+    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.
+
+        Specifying ``update`` and ``apply`` methods is optional and allows certain meta-modules to be used,
+        such as ``tz.m.Online`` or trust regions. Alternatively, simply override the ``step`` method.
+        """
+        return self.step(var)
+
+    def get_H(self, var: Var) -> LinearOperator | None:
+        """returns a ``LinearOperator`` corresponding to hessian or hessian approximation.
+        The hessian approximation is assumed to be for all parameters concatenated to a vector."""
+        # if this method is not defined it searches in children
+        # this should be overwritten to return None if child params are different from this modules params
+        H = None
+        for k,v in self.children.items():
+            H_v = v.get_H(var)
+
+            if (H is not None) and (H_v is not None):
+                raise RuntimeError(f"Two children of {self} have a hessian, second one is {k}={v}")
+
+            if H_v is not None: H = H_v
+
+        return H
 
     def reset(self):
-        """Resets the internal state of the module (e.g. momentum)."""
-        # no complex logic is allowed there because this is overridden by many modules
-        # where super().reset() shouldn't be called
+        """Resets the internal state of the module (e.g. momentum) and all children. By default clears state and global state."""
         self.state.clear()
         self.global_state.clear()
+        for c in self.children.values(): c.reset()
+
+    def reset_for_online(self):
+        """Resets buffers that depend on previous evaluation, such as previous gradient and loss,
+        which may become inaccurate due to mini-batching.
+
+        ``Online`` module calls ``reset_for_online``,
+        then it calls ``update`` with previous parameters,
+        then it calls ``update`` with current parameters,
+        and then ``apply``.
+        """
+        for c in self.children.values(): c.reset_for_online()
 
     def _extra_pack(self):
+        """extra information to store in state_dict of this optimizer.
+        Will be passed to ``_extra_unpack`` when loading the state_dict."""
         return {}
 
     def _extra_unpack(self, x):
-        pass
+        """``_extra_pack`` return will be passed to this method when loading state_dict.
+        This method is called after loading the rest of the state dict"""
+
+
+
+    # ------------------------------ 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,
+    ) -> tuple[Sequence[torch.Tensor], Sequence[torch.Tensor] | None]:
+        """
+        Returns ``(Hvp, rgrad)``, where ``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:
+
+        ```python
+        Hvp, _ = self.hvp(v, at_x0=True, rgrad=None, ..., retain_graph=False)
+        ```
+
+        Multiple samples example:
+
+        ```python
+        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
+
+    def get_generator(self, device: torch.types.Device, seed: int | None):
+        if seed is None: return None
+
+        if 'generator' not in self.global_state:
+            self.global_state['generator'] = torch.Generator(device).manual_seed(seed)
+
+        return self.global_state['generator']
 
 # endregion
 
@@ -440,6 +658,27 @@ def unroll_modules(*modules: Chainable) -> list[Module]:
 
 # region Modular
 # ---------------------------------- Modular --------------------------------- #
+
+class _EvalCounterClosure:
+    """keeps track of how many times closure has been evaluated, and sets closure return"""
+    __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")
+
+        v = self.closure(*args, **kwargs)
+
+        # set closure return on 1st evaluation
+        if self.modular._closure_return is None:
+            self.modular._closure_return = v
+
+        self.modular.num_evaluations += 1
+        return v
+
 # 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):
@@ -456,6 +695,7 @@ class Modular(torch.optim.Optimizer):
     param_groups: list[ChainMap[str, Any]] # pyright:ignore[reportIncompatibleVariableOverride]
 
     def __init__(self, params: Params | torch.nn.Module, *modules: Module):
+        if len(modules) == 0: raise RuntimeError("Empty list of modules passed to `Modular`")
         self.model: torch.nn.Module | None = None
         """The model whose parameters are being optimized, if a model instance was passed to `__init__`."""
         if isinstance(params, torch.nn.Module):
@@ -489,14 +729,33 @@ class Modular(torch.optim.Optimizer):
         for m in self.unrolled_modules: defaults.update(m.defaults)
         super().__init__(param_groups, defaults=defaults)
 
-        # note - this is what super init does:
+        # note - this is what super().__init__(param_groups, defaults=defaults) does:
 
         # self.defaults = defaults
         # for param_group in param_groups:
         #     self.add_param_group(param_group)
 
+        # add_param_group adds a ChainMap where defaults are lowest priority,
+        # and entries specifed in param_groups or scheduler are higher priority.
+        # pytorch schedulers do group["lr"] = new_lr, which sets higher priority key.
+        # in each module, settings passed to that module by calling set_param_groups are highest priority
+
         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)."""
+
+        # reformulations will change the closure to return a different loss (e.g. a sqrt homotopy, gaussian homotopy)
+        # we want to return original loss so this attribute is used
+        self._closure_return = None
+        """on each step, first time a closure is evaluated, this attribute is set to the returned value. `step` method returns this."""
+
+        self.attrs = {}
+        """custom attributes that can be set by modules, for example EMA of weights or best so far"""
+
+        self.should_terminate = False
+        """is set to True by termination criteria modules."""
 
     def add_param_group(self, param_group: dict[str, Any]):
         proc_param_group = _make_param_groups([param_group], differentiable=False)[0]
@@ -542,10 +801,13 @@ class Modular(torch.optim.Optimizer):
 
         id_to_tensor = {state_dict['idx_to_id'][i]: p for i,p in enumerate(state_dict['params'])}
         for m, sd in zip(self.unrolled_modules, state_dict['modules'].values()):
-            m.load_state_dict(sd, id_to_tensor)
+            m._load_state_dict(sd, id_to_tensor)
+
 
+    def step(self, closure=None, loss=None, **kwargs): # pyright: ignore[reportIncompatibleMethodOverride]
+        # clear closure return from previous step
+        self._closure_return = None
 
-    def step(self, closure=None): # pyright: ignore[reportIncompatibleMethodOverride]
         # propagate global per-parameter setting overrides
         for g in self.param_groups:
             settings = dict(g.maps[0]) # ignore defaults
@@ -558,15 +820,17 @@ class Modular(torch.optim.Optimizer):
 
         # create var
         params = [p for g in self.param_groups for p in g['params'] if p.requires_grad]
-        var = Var(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, modular=self, loss=loss, storage=kwargs)
 
         # if closure is None, assume backward has been called and gather grads
         if closure is None:
             var.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+            self.num_evaluations += 1
 
+        n_modules = len(self.modules)
+        if n_modules == 0: raise RuntimeError("There are no modules in this `Modular` optimizer")
         last_module = self.modules[-1]
         last_lr = last_module.defaults.get('lr', None)
-        n_modules = len(self.modules)
 
         # step
         for i, module in enumerate(self.modules):
@@ -586,11 +850,17 @@ class Modular(torch.optim.Optimizer):
             with torch.no_grad():
                 torch._foreach_sub_(params, var.get_update())
 
+        # update attributes
+        self.attrs.update(var.attrs)
+        if var.should_terminate is not None: self.should_terminate = var.should_terminate
+
+        # hooks
         for hook in var.post_step_hooks:
             hook(self, var)
 
         self.current_step += 1
-        return var.loss if var.loss is not None else var.loss_approx
+        #return var.loss if var.loss is not None else var.loss_approx
+        return self._closure_return
 
     def __repr__(self):
         return f'Modular({", ".join(str(m) for m in self.modules)})'
@@ -606,6 +876,21 @@ class Chain(Module):
         for i, module in enumerate(flat_modules):
             self.set_child(f'module_{i}', module)
 
+    def update(self, var):
+        # note here that `update` and `apply` shouldn't be used directly
+        # as it will update all modules, and then apply all modules
+        # it is used in specific cases like Chain as trust region hessian module
+        for i in range(len(self.children)):
+            self.children[f'module_{i}'].update(var)
+            if var.stop: break
+        return var
+
+    def apply(self, var):
+        for i in range(len(self.children)):
+            var = self.children[f'module_{i}'].apply(var)
+            if var.stop: break
+        return var
+
     def step(self, var):
         for i in range(len(self.children)):
             var = self.children[f'module_{i}'].step(var)
@@ -616,7 +901,7 @@ class Chain(Module):
         s = self.__class__.__name__
         if self.children:
             if s == 'Chain': s = 'C' # to shorten it
-            s = f'{s}({", ".join(str(m) for m in self.children.values())}'
+            s = f'{s}({", ".join(str(m) for m in self.children.values())})'
         return s
 
 def maybe_chain(*modules: Chainable) -> Module: