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

docs/source/conf.py

@@ -6,10 +6,10 @@
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 import sys, os
-#sys.path.insert(0, os.path.abspath('.../src'))  
+#sys.path.insert(0, os.path.abspath('.../src'))
 
 project = 'torchzero'
-copyright = '2024, Ivan Nikishev'
+copyright = '2025, Ivan Nikishev'
 author = 'Ivan Nikishev'
 
 # -- General configuration ---------------------------------------------------
@@ -24,10 +24,12 @@ extensions = [
     'sphinx.ext.githubpages',
     'sphinx.ext.napoleon',
     'autoapi.extension',
+    "myst_nb",
+
     # 'sphinx_rtd_theme',
 ]
 autosummary_generate = True
-autoapi_dirs = ['../../src']
+autoapi_dirs = ['../../torchzero']
 autoapi_type = "python"
 # autoapi_ignore = ["*/tensorlist.py"]
 
@@ -48,7 +50,7 @@ exclude_patterns = []
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
 #html_theme = 'alabaster'
-html_theme = 'furo'
+html_theme = 'sphinx_rtd_theme'
 html_static_path = ['_static']
 
 

docs/source/docstring template.py

@@ -0,0 +1,46 @@
+class MyModule:
+    """[One-line summary of the class].
+
+    [A more detailed description of the class, explaining its purpose, how it
+    works, and its typical use cases. You can use multiple paragraphs.]
+
+    .. note::
+        [Optional: Add important notes, warnings, or usage guidelines here.
+        For example, you could mention if a closure is required, discuss
+        stability, or highlight performance characteristics. Use the `.. note::`
+        directive to make it stand out in the documentation.]
+
+    Args:
+        param1 (type, optional):
+            [Description of the first parameter. Use :code:`backticks` for
+            inline code like variable names or specific values like ``"autograd"``.
+            Explain what the parameter does.] Defaults to [value].
+        param2 (type):
+            [Description of a mandatory parameter (no "optional" or "Defaults to").]
+        **kwargs:
+            [If you accept keyword arguments, describe what they are used for.]
+
+    Examples:
+        [A title or short sentence describing the first example]:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                ...
+            )
+
+        [A title or short sentence for a second, different example]:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                ...
+            )
+
+    References:
+        - [Optional: A citation for a relevant paper, book, or algorithm.]
+        - [Optional: A link to a blog post or website with more information.]
+
+    """

tests/test_identical.py

@@ -96,8 +96,7 @@ def _assert_identical_device(opt_fn: Callable, merge: bool, use_closure: bool, s
 
 @pytest.mark.parametrize('amsgrad', [True, False])
 def test_adam(amsgrad):
-    # torch_fn = lambda p: torch.optim.Adam(p, lr=1, amsgrad=amsgrad)
-    # pytorch applies debiasing separately so it is applied before epsilo
+    torch_fn = lambda p: torch.optim.Adam(p, lr=1, amsgrad=amsgrad)
     tz_fn = lambda p: tz.Modular(p, tz.m.Adam(amsgrad=amsgrad))
     tz_fn2 = lambda p: tz.Modular(p, tz.m.Adam(amsgrad=amsgrad), tz.m.LR(1)) # test LR fusing
     tz_fn3 = lambda p: tz.Modular(p, tz.m.Adam(amsgrad=amsgrad), tz.m.LR(1), tz.m.Add(1), tz.m.Sub(1))
@@ -133,7 +132,7 @@ def test_adam(amsgrad):
                 tz.m.Debias2(beta=0.999),
                 tz.m.Add(1e-8)]
         ))
-    tz_fns = (tz_fn, tz_fn2, tz_fn3, tz_fn4, tz_fn5, tz_fn_ops, tz_fn_ops2, tz_fn_ops3, tz_fn_ops4)
+    tz_fns = (torch_fn, tz_fn, tz_fn2, tz_fn3, tz_fn4, tz_fn5, tz_fn_ops, tz_fn_ops2, tz_fn_ops3, tz_fn_ops4)
 
     _assert_identical_opts(tz_fns, merge=True, use_closure=True, device='cpu', steps=10)
     for fn in tz_fns:

tests/test_opts.py

@@ -292,42 +292,42 @@ FDM_central2 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='central2'), tz.m.LR(0.1)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.FDM(), tz.m.LR(0.1)),
     needs_closure=True,
-    func='booth', steps=50, loss=1e-7, merge_invariant=True,
+    func='booth', steps=50, loss=1e-6, merge_invariant=True,
     sphere_steps=2, sphere_loss=340,
 )
 FDM_forward2 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='forward2'), tz.m.LR(0.1)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='forward2'), tz.m.LR(0.1)),
     needs_closure=True,
-    func='booth', steps=50, loss=1e-7, merge_invariant=True,
+    func='booth', steps=50, loss=1e-6, merge_invariant=True,
     sphere_steps=2, sphere_loss=340,
 )
 FDM_backward2 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='backward2'), tz.m.LR(0.1)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='backward2'), tz.m.LR(0.1)),
     needs_closure=True,
-    func='booth', steps=50, loss=2e-7, merge_invariant=True,
+    func='booth', steps=50, loss=1e-6, merge_invariant=True,
     sphere_steps=2, sphere_loss=340,
 )
 FDM_forward3 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='forward3'), tz.m.LR(0.1)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='forward3'), tz.m.LR(0.1)),
     needs_closure=True,
-    func='booth', steps=50, loss=3e-7, merge_invariant=True,
+    func='booth', steps=50, loss=1e-6, merge_invariant=True,
     sphere_steps=2, sphere_loss=340,
 )
 FDM_backward3 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='backward3'), tz.m.LR(0.1)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='backward3'), tz.m.LR(0.1)),
     needs_closure=True,
-    func='booth', steps=50, loss=3e-7, merge_invariant=True,
+    func='booth', steps=50, loss=1e-6, merge_invariant=True,
     sphere_steps=2, sphere_loss=340,
 )
 FDM_central4 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='central4'), tz.m.LR(0.1)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.FDM(formula='central4'), tz.m.LR(0.1)),
     needs_closure=True,
-    func='booth', steps=50, loss=2e-8, merge_invariant=True,
+    func='booth', steps=50, loss=1e-6, merge_invariant=True,
     sphere_steps=2, sphere_loss=340,
 )
 
@@ -466,8 +466,8 @@ AdaptiveBacktracking = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.AdaptiveBacktracking()),
     sphere_opt=lambda p: tz.Modular(p, tz.m.AdaptiveBacktracking()),
     needs_closure=True,
-    func='booth', steps=50, loss=0, merge_invariant=True,
-    sphere_steps=2, sphere_loss=0,
+    func='booth', steps=50, loss=1e-12, merge_invariant=True,
+    sphere_steps=2, sphere_loss=1e-10,
 )
 AdaptiveBacktracking_try_negative = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.Mul(-1), tz.m.AdaptiveBacktracking(try_negative=True)),
@@ -494,15 +494,6 @@ StrongWolfe = Run(
     sphere_steps=2, sphere_loss=0,
 )
 
-# ------------------------- line_search/trust_region ------------------------- #
-TrustRegion = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.TrustRegion()),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.TrustRegion(init=0.1)),
-    needs_closure=True,
-    func='booth', steps=50, loss=0.1, merge_invariant=True,
-    sphere_steps=10, sphere_loss=1e-5,
-)
-
 # ----------------------------------- lr/lr ---------------------------------- #
 LR = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.LR(0.1)),
@@ -587,8 +578,8 @@ UpdateGradientSignConsistency = Run(
     sphere_steps=10, sphere_loss=2,
 )
 IntermoduleCautious = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.IntermoduleCautious(tz.m.NAG(), tz.m.BFGS()), tz.m.LR(0.01)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.IntermoduleCautious(tz.m.NAG(), tz.m.BFGS()), tz.m.LR(0.1)),
+    func_opt=lambda p: tz.Modular(p, tz.m.IntermoduleCautious(tz.m.NAG(), tz.m.BFGS(ptol_reset=True)), tz.m.LR(0.01)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.IntermoduleCautious(tz.m.NAG(), tz.m.BFGS(ptol_reset=True)), tz.m.LR(0.1)),
     needs_closure=False,
     func='booth', steps=50, loss=1e-4, merge_invariant=True,
     sphere_steps=10, sphere_loss=0.1,
@@ -601,8 +592,8 @@ ScaleByGradCosineSimilarity = Run(
     sphere_steps=10, sphere_loss=0.1,
 )
 ScaleModulesByCosineSimilarity = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.ScaleModulesByCosineSimilarity(tz.m.HeavyBall(0.9), tz.m.BFGS()),tz.m.LR(0.05)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.ScaleModulesByCosineSimilarity(tz.m.HeavyBall(0.9), tz.m.BFGS()),tz.m.LR(0.1)),
+    func_opt=lambda p: tz.Modular(p, tz.m.ScaleModulesByCosineSimilarity(tz.m.HeavyBall(0.9), tz.m.BFGS(ptol_reset=True)),tz.m.LR(0.05)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.ScaleModulesByCosineSimilarity(tz.m.HeavyBall(0.9), tz.m.BFGS(ptol_reset=True)),tz.m.LR(0.1)),
     needs_closure=False,
     func='booth', steps=50, loss=0.005, merge_invariant=True,
     sphere_steps=10, sphere_loss=0.1,
@@ -684,8 +675,8 @@ GradAccumulation = Run(
     sphere_steps=20, sphere_loss=1e-11,
 )
 NegateOnLossIncrease = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.HeavyBall(), tz.m.LR(0.02), tz.m.NegateOnLossIncrease(),),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.HeavyBall(), tz.m.LR(0.1), tz.m.NegateOnLossIncrease(),),
+    func_opt=lambda p: tz.Modular(p, tz.m.HeavyBall(), tz.m.LR(0.02), tz.m.NegateOnLossIncrease(True),),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.HeavyBall(), tz.m.LR(0.1), tz.m.NegateOnLossIncrease(True),),
     needs_closure=True,
     func='booth', steps=50, loss=0.1, merge_invariant=True,
     sphere_steps=20, sphere_loss=0.001,
@@ -693,7 +684,7 @@ NegateOnLossIncrease = Run(
 # -------------------------------- misc/switch ------------------------------- #
 Alternate = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.Alternate(tz.m.Adagrad(), tz.m.Adam(), tz.m.RMSprop()), tz.m.LR(1)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.Alternate(tz.m.Adagrad(), tz.m.Adam(), tz.m.RMSprop()), tz.m.LR(1)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.Alternate(tz.m.Adagrad(), tz.m.Adam(), tz.m.RMSprop()), tz.m.LR(0.1)),
     needs_closure=False,
     func='booth', steps=50, loss=1, merge_invariant=True,
     sphere_steps=20, sphere_loss=20,
@@ -734,24 +725,24 @@ Shampoo = Run(
 
 # ------------------------- quasi_newton/quasi_newton ------------------------ #
 BFGS = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.BFGS(), tz.m.StrongWolfe()),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.BFGS(), tz.m.StrongWolfe()),
+    func_opt=lambda p: tz.Modular(p, tz.m.BFGS(ptol_reset=True), tz.m.StrongWolfe()),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.BFGS(ptol_reset=True), tz.m.StrongWolfe()),
     needs_closure=True,
-    func='rosen', steps=50, loss=0, merge_invariant=True,
-    sphere_steps=10, sphere_loss=0,
+    func='rosen', steps=50, loss=1e-10, merge_invariant=True,
+    sphere_steps=10, sphere_loss=1e-10,
 )
 SR1 = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.SR1(), tz.m.StrongWolfe()),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.SR1(), tz.m.StrongWolfe()),
+    func_opt=lambda p: tz.Modular(p, tz.m.SR1(ptol_reset=True), tz.m.StrongWolfe()),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.SR1(ptol_reset=True), tz.m.StrongWolfe()),
     needs_closure=True,
     func='rosen', steps=50, loss=1e-12, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 SSVM = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.SSVM(1), tz.m.StrongWolfe()),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.SSVM(1), tz.m.StrongWolfe()),
+    func_opt=lambda p: tz.Modular(p, tz.m.SSVM(1, ptol_reset=True), tz.m.StrongWolfe()),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.SSVM(1, ptol_reset=True), tz.m.StrongWolfe()),
     needs_closure=True,
-    func='rosen', steps=50, loss=1e-10, merge_invariant=True,
+    func='rosen', steps=50, loss=0.5, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 
@@ -766,21 +757,21 @@ LBFGS = Run(
 
 # ----------------------------- quasi_newton/lsr1 ---------------------------- #
 LSR1 = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.LSR1(), tz.m.StrongWolfe()),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.LSR1(), tz.m.StrongWolfe()),
+    func_opt=lambda p: tz.Modular(p, tz.m.LSR1(scale_second=True), tz.m.StrongWolfe()),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.LSR1(scale_second=True), tz.m.StrongWolfe()),
     needs_closure=True,
     func='rosen', steps=50, loss=0, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 
-# ---------------------------- quasi_newton/olbfgs --------------------------- #
-OnlineLBFGS = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.OnlineLBFGS(), tz.m.StrongWolfe()),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.OnlineLBFGS(), tz.m.StrongWolfe()),
-    needs_closure=True,
-    func='rosen', steps=50, loss=0, merge_invariant=True,
-    sphere_steps=10, sphere_loss=0,
-)
+# # ---------------------------- quasi_newton/olbfgs --------------------------- #
+# OnlineLBFGS = Run(
+#     func_opt=lambda p: tz.Modular(p, tz.m.OnlineLBFGS(), tz.m.StrongWolfe()),
+#     sphere_opt=lambda p: tz.Modular(p, tz.m.OnlineLBFGS(), tz.m.StrongWolfe()),
+#     needs_closure=True,
+#     func='rosen', steps=50, loss=0, merge_invariant=True,
+#     sphere_steps=10, sphere_loss=0,
+# )
 
 # ---------------------------- second_order/newton --------------------------- #
 Newton = Run(
@@ -802,8 +793,8 @@ NewtonCG = Run(
 
 # ---------------------------- smoothing/gaussian ---------------------------- #
 GaussianHomotopy = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.GaussianHomotopy(10, 1, tol=1e-1, seed=0), tz.m.BFGS(), tz.m.StrongWolfe()),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.GaussianHomotopy(10, 1, tol=1e-1, seed=0), tz.m.BFGS(), tz.m.StrongWolfe()),
+    func_opt=lambda p: tz.Modular(p, tz.m.GaussianHomotopy(10, 1, tol=1e-1, seed=0), tz.m.BFGS(ptol_reset=True), tz.m.StrongWolfe()),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.GaussianHomotopy(10, 1, tol=1e-1, seed=0), tz.m.BFGS(ptol_reset=True), tz.m.StrongWolfe()),
     needs_closure=True,
     func='booth', steps=20, loss=0.1, merge_invariant=True,
     sphere_steps=10, sphere_loss=200,
@@ -860,7 +851,7 @@ SophiaH = Run(
     sphere_steps=10, sphere_loss=40,
 )
 
-# -------------------------- optimizers/higher_order ------------------------- #
+# -------------------------- higher_order ------------------------- #
 HigherOrderNewton = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.HigherOrderNewton(trust_method=None)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.HigherOrderNewton(2, trust_method=None)),
@@ -869,6 +860,24 @@ HigherOrderNewton = Run(
     sphere_steps=1, sphere_loss=1e-10,
 )
 
+# ---------------------------- optimizers/ladagrad --------------------------- #
+LMAdagrad = Run(
+    func_opt=lambda p: tz.Modular(p, tz.m.LMAdagrad(), tz.m.LR(4)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.LMAdagrad(), tz.m.LR(5)),
+    needs_closure=False,
+    func='booth', steps=50, loss=1e-6, merge_invariant=True,
+    sphere_steps=20, sphere_loss=1e-9,
+)
+
+# ------------------------------ optimizers/adan ----------------------------- #
+Adan = Run(
+    func_opt=lambda p: tz.Modular(p, tz.m.Adan(), tz.m.LR(1)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.Adan(), tz.m.LR(0.1)),
+    needs_closure=False,
+    func='booth', steps=50, loss=60, merge_invariant=True,
+    sphere_steps=20, sphere_loss=60,
+)
+
 # ------------------------------------ CGs ----------------------------------- #
 for CG in (tz.m.PolakRibiere, tz.m.FletcherReeves, tz.m.HestenesStiefel, tz.m.DaiYuan, tz.m.LiuStorey, tz.m.ConjugateDescent, tz.m.HagerZhang, tz.m.HybridHS_DY, tz.m.ProjectedGradientMethod):
     for func_steps,sphere_steps_ in ([3,2], [10,10]): # CG should converge on 2D quadratic after 2nd step
@@ -885,17 +894,22 @@ for CG in (tz.m.PolakRibiere, tz.m.FletcherReeves, tz.m.HestenesStiefel, tz.m.Da
 # stability test
 for QN in (
     tz.m.BFGS,
+    partial(tz.m.BFGS, inverse=False),
     tz.m.SR1,
+    partial(tz.m.SR1, inverse=False),
     tz.m.DFP,
+    partial(tz.m.DFP, inverse=False),
     tz.m.BroydenGood,
+    partial(tz.m.BroydenGood, inverse=False),
     tz.m.BroydenBad,
+    partial(tz.m.BroydenBad, inverse=False),
     tz.m.Greenstadt1,
     tz.m.Greenstadt2,
-    tz.m.ColumnUpdatingMethod,
+    tz.m.ICUM,
     tz.m.ThomasOptimalMethod,
     tz.m.FletcherVMM,
     tz.m.Horisho,
-    lambda scale_first: tz.m.Horisho(scale_first=scale_first, inner=tz.m.GradientCorrection()),
+    partial(tz.m.Horisho, inner=tz.m.GradientCorrection()),
     tz.m.Pearson,
     tz.m.ProjectedNewtonRaphson,
     tz.m.PSB,
@@ -903,8 +917,8 @@ for QN in (
     tz.m.SSVM,
 ):
     Run(
-        func_opt=lambda p: tz.Modular(p, QN(scale_first=False), tz.m.StrongWolfe()),
-        sphere_opt=lambda p: tz.Modular(p, QN(scale_first=False), tz.m.StrongWolfe()),
+        func_opt=lambda p: tz.Modular(p, QN(scale_first=False, ptol_reset=True), tz.m.StrongWolfe()),
+        sphere_opt=lambda p: tz.Modular(p, QN(scale_first=False, ptol_reset=True), tz.m.StrongWolfe()),
         needs_closure=True,
         func='lstsq', steps=50, loss=1e-10, merge_invariant=False,
         sphere_steps=10, sphere_loss=1e-20,

tests/test_vars.py

@@ -156,6 +156,7 @@ 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 = }'

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
 
 
@@ -109,6 +110,9 @@ class Var:
         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:`var.loss`.
         Do not call this at perturbed parameters. Backward always zeroes grads before recomputing."""
@@ -186,6 +190,7 @@ class Var:
         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:
@@ -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, 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 ::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]
@@ -558,11 +689,12 @@ 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)
 
         # 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
 
         last_module = self.modules[-1]
         last_lr = last_module.defaults.get('lr', None)