torchzero 0.3.9__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

@@ -1,4 +1,9 @@
-"""snity tests to make sure everything works and converges on basic functions"""
+"""
+Sanity tests to make sure everything works.
+
+This will show major convergence regressions, but that is not the main purpose. Mainly this makes sure modules
+don't error or become unhinged with different parameter shapes.
+"""
 from collections.abc import Callable
 from functools import partial
 
@@ -68,6 +73,7 @@ def _run_objective(opt: tz.Modular, objective: Callable, use_closure: bool, step
             assert torch.isfinite(loss), f"{opt}: Inifinite loss - {[l.item() for l in losses]}"
             losses.append(loss)
 
+    losses.append(objective())
     return torch.stack(losses).nan_to_num(0,10000,10000).min()
 
 def _run_func(opt_fn: Callable, func:str, merge: bool, use_closure: bool, steps: int):
@@ -286,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,
 )
 
@@ -460,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)),
@@ -488,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)),
@@ -524,7 +521,7 @@ PolyakStepSize = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.PolyakStepSize()),
     sphere_opt=lambda p: tz.Modular(p, tz.m.PolyakStepSize()),
     needs_closure=True,
-    func='booth', steps=50, loss=1e-11, merge_invariant=True,
+    func='booth', steps=50, loss=1e-7, merge_invariant=True,
     sphere_steps=10, sphere_loss=0.002,
 )
 RandomStepSize = Run(
@@ -581,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,
@@ -595,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,
@@ -604,44 +601,44 @@ ScaleModulesByCosineSimilarity = Run(
 
 # ------------------------- momentum/matrix_momentum ------------------------- #
 MatrixMomentum_forward = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_mode='forward'), tz.m.LR(0.01)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_mode='forward'), tz.m.LR(0.5)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_method='forward'), tz.m.LR(0.01)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_method='forward'), tz.m.LR(0.5)),
     needs_closure=True,
     func='booth', steps=50, loss=0.05, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 MatrixMomentum_forward = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_mode='central'), tz.m.LR(0.01)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_mode='central'), tz.m.LR(0.5)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_method='central'), tz.m.LR(0.01)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_method='central'), tz.m.LR(0.5)),
     needs_closure=True,
     func='booth', steps=50, loss=0.05, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 MatrixMomentum_forward = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_mode='autograd'), tz.m.LR(0.01)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_mode='autograd'), tz.m.LR(0.5)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_method='autograd'), tz.m.LR(0.01)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(hvp_method='autograd'), tz.m.LR(0.5)),
     needs_closure=True,
     func='booth', steps=50, loss=0.05, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 
 AdaptiveMatrixMomentum_forward = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_mode='forward'), tz.m.LR(0.05)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_mode='forward'), tz.m.LR(0.5)),
+    func_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_method='forward'), tz.m.LR(0.05)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_method='forward'), tz.m.LR(0.5)),
     needs_closure=True,
     func='booth', steps=50, loss=0.002, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 AdaptiveMatrixMomentum_central = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_mode='central'), tz.m.LR(0.05)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_mode='central'), tz.m.LR(0.5)),
+    func_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_method='central'), tz.m.LR(0.05)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_method='central'), tz.m.LR(0.5)),
     needs_closure=True,
     func='booth', steps=50, loss=0.002, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 AdaptiveMatrixMomentum_autograd = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_mode='autograd'), tz.m.LR(0.05)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_mode='autograd'), tz.m.LR(0.5)),
+    func_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_method='autograd'), tz.m.LR(0.05)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.AdaptiveMatrixMomentum(hvp_method='autograd'), tz.m.LR(0.5)),
     needs_closure=True,
     func='booth', steps=50, loss=0.002, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
@@ -678,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,
@@ -687,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,
@@ -719,33 +716,33 @@ Lion = Run(
 )
 # ---------------------------- optimizers/shampoo ---------------------------- #
 Shampoo = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.GraftModules(tz.m.Shampoo(), tz.m.RMSprop()), tz.m.LR(0.1)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.GraftModules(tz.m.Shampoo(), tz.m.RMSprop()), tz.m.LR(0.2)),
+    func_opt=lambda p: tz.Modular(p, tz.m.GraftModules(tz.m.Shampoo(), tz.m.RMSprop()), tz.m.LR(4)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.GraftModules(tz.m.Shampoo(), tz.m.RMSprop()), tz.m.LR(0.1)),
     needs_closure=False,
-    func='booth', steps=50, loss=200, merge_invariant=False,
-    sphere_steps=20, sphere_loss=1e-3, # merge and unmerge lrs are very different so need to test convergence separately somewhere
+    func='booth', steps=50, loss=0.02, merge_invariant=False,
+    sphere_steps=20, sphere_loss=1, # merge and unmerge lrs are very different so need to test convergence separately somewhere
 )
 
 # ------------------------- 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,
 )
 
@@ -760,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(
@@ -791,13 +788,13 @@ NewtonCG = Run(
     sphere_opt=lambda p: tz.Modular(p, tz.m.NewtonCG(), tz.m.StrongWolfe()),
     needs_closure=True,
     func='rosen', steps=20, loss=1e-7, merge_invariant=True,
-    sphere_steps=2, sphere_loss=1e-6,
+    sphere_steps=2, sphere_loss=3e-4,
 )
 
 # ---------------------------- 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,
@@ -854,8 +851,35 @@ SophiaH = Run(
     sphere_steps=10, sphere_loss=40,
 )
 
+# -------------------------- 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)),
+    needs_closure=True,
+    func='rosen', steps=1, loss=2e-10, merge_invariant=True,
+    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):
+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
         # but also test 10 to make sure it doesn't explode after converging
         Run(
@@ -868,10 +892,33 @@ for CG in (tz.m.PolakRibiere, tz.m.FletcherReeves, tz.m.HestenesStiefel, tz.m.Da
 
 # ------------------------------- QN stability ------------------------------- #
 # stability test
-for QN in (tz.m.BFGS, tz.m.SR1, tz.m.DFP, tz.m.BroydenGood, tz.m.BroydenBad, tz.m.Greenstadt1, tz.m.Greenstadt2, tz.m.ColumnUpdatingMethod,  tz.m.ThomasOptimalMethod, tz.m.PSB, tz.m.Pearson2, tz.m.SSVM):
+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.ICUM,
+    tz.m.ThomasOptimalMethod,
+    tz.m.FletcherVMM,
+    tz.m.Horisho,
+    partial(tz.m.Horisho, inner=tz.m.GradientCorrection()),
+    tz.m.Pearson,
+    tz.m.ProjectedNewtonRaphson,
+    tz.m.PSB,
+    tz.m.McCormick,
+    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_tensorlist.py

@@ -1261,8 +1261,8 @@ def test_reduction_ops(simple_tl: TensorList, reduction_method, dim, keepdim):
         elif reduction_method == 'quantile': expected = vec.quantile(q)
         else:
             pytest.fail("Unknown global reduction")
-            assert False, 'sus'
-        assert torch.allclose(result, expected)
+            assert False, reduction_method
+        assert torch.allclose(result, expected, atol=1e-4)
     else:
         expected_list = []
         for t in simple_tl: