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

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

@@ -56,14 +56,17 @@ def _run_objective(opt: tz.Modular, objective: Callable, use_closure: bool, step
         if use_closure:
             def closure(backward=True):
                 loss = objective()
+                losses.append(loss.detach())
                 if backward:
                     opt.zero_grad()
                     loss.backward()
                 return loss
-            loss = opt.step(closure)
-            assert loss is not None
-            assert torch.isfinite(loss), f"{opt}: Inifinite loss - {[l.item() for l in losses]}"
-            losses.append(loss)
+            ret = opt.step(closure)
+            assert ret is not None # the return should be the loss
+            with torch.no_grad():
+                loss = objective() # in case f(x_0) is not evaluated
+                assert torch.isfinite(loss), f"{opt}: Inifinite loss - {[l.item() for l in losses]}"
+                losses.append(loss.detach())
 
         else:
             loss = objective()
@@ -71,7 +74,7 @@ def _run_objective(opt: tz.Modular, objective: Callable, use_closure: bool, step
             loss.backward()
             opt.step()
             assert torch.isfinite(loss), f"{opt}: Inifinite loss - {[l.item() for l in losses]}"
-            losses.append(loss)
+            losses.append(loss.detach())
 
     losses.append(objective())
     return torch.stack(losses).nan_to_num(0,10000,10000).min()
@@ -292,42 +295,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,
 )
 
@@ -374,6 +377,21 @@ RandomizedFDM_central4 = Run(
     func='booth', steps=50, loss=10, merge_invariant=True,
     sphere_steps=100, sphere_loss=450,
 )
+RandomizedFDM_forward4 = Run(
+    func_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='forward4', seed=0), tz.m.LR(0.01)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='forward4', seed=0), tz.m.LR(0.001)),
+    needs_closure=True,
+    func='booth', steps=50, loss=10, merge_invariant=True,
+    sphere_steps=100, sphere_loss=450,
+)
+RandomizedFDM_forward5 = Run(
+    func_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='forward5', seed=0), tz.m.LR(0.01)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='forward5', seed=0), tz.m.LR(0.001)),
+    needs_closure=True,
+    func='booth', steps=50, loss=10, merge_invariant=True,
+    sphere_steps=100, sphere_loss=450,
+)
+
 
 RandomizedFDM_4samples = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(n_samples=4, seed=0), tz.m.LR(0.1)),
@@ -455,25 +473,11 @@ Backtracking = Run(
     func='booth', steps=50, loss=0, merge_invariant=True,
     sphere_steps=2, sphere_loss=0,
 )
-Backtracking_try_negative = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.Mul(-1), tz.m.Backtracking(try_negative=True)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.Mul(-1), tz.m.Backtracking(try_negative=True)),
-    needs_closure=True,
-    func='booth', steps=50, loss=1e-9, merge_invariant=True,
-    sphere_steps=2, sphere_loss=1e-10,
-)
 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,
-)
-AdaptiveBacktracking_try_negative = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.Mul(-1), tz.m.AdaptiveBacktracking(try_negative=True)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.Mul(-1), tz.m.AdaptiveBacktracking(try_negative=True)),
-    needs_closure=True,
-    func='booth', steps=50, loss=1e-8, merge_invariant=True,
+    func='booth', steps=50, loss=1e-11, merge_invariant=True,
     sphere_steps=2, sphere_loss=1e-10,
 )
 # ----------------------------- line_search/scipy ---------------------------- #
@@ -494,15 +498,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 +582,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_restart=True)), tz.m.LR(0.01)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.IntermoduleCautious(tz.m.NAG(), tz.m.BFGS(ptol_restart=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 +596,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_restart=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_restart=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,
@@ -610,47 +605,69 @@ ScaleModulesByCosineSimilarity = Run(
 
 # ------------------------- momentum/matrix_momentum ------------------------- #
 MatrixMomentum_forward = Run(
-    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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.01, hvp_method='forward'),),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='forward')),
     needs_closure=True,
     func='booth', steps=50, loss=0.05, merge_invariant=True,
-    sphere_steps=10, sphere_loss=0,
+    sphere_steps=10, sphere_loss=0.01,
 )
 MatrixMomentum_forward = Run(
-    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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.01, hvp_method='central')),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='central')),
     needs_closure=True,
     func='booth', steps=50, loss=0.05, merge_invariant=True,
-    sphere_steps=10, sphere_loss=0,
+    sphere_steps=10, sphere_loss=0.01,
 )
 MatrixMomentum_forward = Run(
-    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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.01, hvp_method='autograd')),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='autograd')),
     needs_closure=True,
     func='booth', steps=50, loss=0.05, merge_invariant=True,
-    sphere_steps=10, sphere_loss=0,
+    sphere_steps=10, sphere_loss=0.01,
 )
 
 AdaptiveMatrixMomentum_forward = Run(
-    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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='forward', adaptive=True)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='forward', adaptive=True)),
     needs_closure=True,
-    func='booth', steps=50, loss=0.002, merge_invariant=True,
-    sphere_steps=10, sphere_loss=0,
+    func='booth', steps=50, loss=0.05, merge_invariant=True,
+    sphere_steps=10, sphere_loss=0.05,
 )
 AdaptiveMatrixMomentum_central = Run(
-    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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='central', adaptive=True)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='central', adaptive=True)),
     needs_closure=True,
-    func='booth', steps=50, loss=0.002, merge_invariant=True,
-    sphere_steps=10, sphere_loss=0,
+    func='booth', steps=50, loss=0.05, merge_invariant=True,
+    sphere_steps=10, sphere_loss=0.05,
 )
 AdaptiveMatrixMomentum_autograd = Run(
-    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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='autograd', adaptive=True)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='autograd', adaptive=True)),
     needs_closure=True,
-    func='booth', steps=50, loss=0.002, merge_invariant=True,
-    sphere_steps=10, sphere_loss=0,
+    func='booth', steps=50, loss=0.05, merge_invariant=True,
+    sphere_steps=10, sphere_loss=0.05,
+)
+
+StochasticAdaptiveMatrixMomentum_forward = Run(
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='forward', adaptive=True, adapt_freq=1)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='forward', adaptive=True, adapt_freq=1)),
+    needs_closure=True,
+    func='booth', steps=50, loss=0.05, merge_invariant=True,
+    sphere_steps=10, sphere_loss=0.05,
+)
+StochasticAdaptiveMatrixMomentum_central = Run(
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='central', adaptive=True, adapt_freq=1)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='central', adaptive=True, adapt_freq=1)),
+    needs_closure=True,
+    func='booth', steps=50, loss=0.05, merge_invariant=True,
+    sphere_steps=10, sphere_loss=0.05,
+)
+StochasticAdaptiveMatrixMomentum_autograd = Run(
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='autograd', adaptive=True, adapt_freq=1)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='autograd', adaptive=True, adapt_freq=1)),
+    needs_closure=True,
+    func='booth', steps=50, loss=0.05, merge_invariant=True,
+    sphere_steps=10, sphere_loss=0.05,
 )
 
 # EMA, momentum are covered by test_identical
@@ -677,15 +694,15 @@ UpdateSign = Run(
     sphere_steps=10, sphere_loss=0,
 )
 GradAccumulation = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.GradientAccumulation(tz.m.LR(0.05), 10), ),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.GradientAccumulation(tz.m.LR(0.5), 10), ),
+    func_opt=lambda p: tz.Modular(p, tz.m.GradientAccumulation(n=10), tz.m.LR(0.05)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.GradientAccumulation(n=10), tz.m.LR(0.5)),
     needs_closure=False,
     func='booth', steps=50, loss=25, merge_invariant=True,
     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 +710,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 +751,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_restart=True), tz.m.StrongWolfe()),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.BFGS(ptol_restart=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_restart=True, scale_first=True), tz.m.StrongWolfe(fallback=False)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.SR1(scale_first=True), tz.m.StrongWolfe(fallback=False)),
     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_restart=True), tz.m.StrongWolfe(fallback=True)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.SSVM(1, ptol_restart=True), tz.m.StrongWolfe(fallback=True)),
     needs_closure=True,
-    func='rosen', steps=50, loss=1e-10, merge_invariant=True,
+    func='rosen', steps=50, loss=0.2, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 
@@ -766,26 +783,26 @@ 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(), tz.m.StrongWolfe(c2=0.1, fallback=True)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.LSR1(), tz.m.StrongWolfe(c2=0.1, fallback=True)),
     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(
-    func_opt=lambda p: tz.Modular(p, tz.m.Newton(), tz.m.StrongWolfe()),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.Newton(), tz.m.StrongWolfe()),
+    func_opt=lambda p: tz.Modular(p, tz.m.Newton(), tz.m.StrongWolfe(fallback=True)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.Newton(), tz.m.StrongWolfe(fallback=True)),
     needs_closure=True,
     func='rosen', steps=20, loss=1e-7, merge_invariant=True,
     sphere_steps=2, sphere_loss=1e-9,
@@ -793,8 +810,8 @@ Newton = Run(
 
 # --------------------------- second_order/newton_cg -------------------------- #
 NewtonCG = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.NewtonCG(), tz.m.StrongWolfe()),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.NewtonCG(), tz.m.StrongWolfe()),
+    func_opt=lambda p: tz.Modular(p, tz.m.NewtonCG(), tz.m.StrongWolfe(fallback=True)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.NewtonCG(), tz.m.StrongWolfe(fallback=True)),
     needs_closure=True,
     func='rosen', steps=20, loss=1e-7, merge_invariant=True,
     sphere_steps=2, sphere_loss=3e-4,
@@ -802,11 +819,11 @@ 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.GradientSampling([tz.m.BFGS(), tz.m.Backtracking()], 1, 10, termination=tz.m.TerminateByUpdateNorm(1e-1), seed=0)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.GradientSampling([tz.m.BFGS(), tz.m.Backtracking()], 1e-1, 10, termination=tz.m.TerminateByUpdateNorm(1e-1), seed=0)),
     needs_closure=True,
-    func='booth', steps=20, loss=0.1, merge_invariant=True,
-    sphere_steps=10, sphere_loss=200,
+    func='booth', steps=20, loss=0.01, merge_invariant=True,
+    sphere_steps=10, sphere_loss=1,
 )
 
 # ---------------------------- smoothing/laplacian --------------------------- #
@@ -860,7 +877,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,15 +886,33 @@ 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 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.DYHS, 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(
             func_opt=lambda p: tz.Modular(p, CG(), tz.m.StrongWolfe(c2=0.1)),
             sphere_opt=lambda p: tz.Modular(p, CG(), tz.m.StrongWolfe(c2=0.1)),
             needs_closure=True,
-            func='lstsq', steps=func_steps, loss=1e-10, merge_invariant=False, # strong wolfe adds float imprecision
+            func='lstsq', steps=func_steps, loss=1e-10, merge_invariant=True,
             sphere_steps=sphere_steps_, sphere_loss=0,
         )
 
@@ -885,17 +920,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,10 +943,10 @@ 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_restart=True), tz.m.StrongWolfe()),
+        sphere_opt=lambda p: tz.Modular(p, QN(scale_first=False, ptol_restart=True), tz.m.StrongWolfe()),
         needs_closure=True,
-        func='lstsq', steps=50, loss=1e-10, merge_invariant=False,
+        func='lstsq', steps=50, loss=1e-10, merge_invariant=True,
         sphere_steps=10, sphere_loss=1e-20,
     )
 

tests/test_tensorlist.py

@@ -977,22 +977,23 @@ def test_rademacher_like(big_tl: TensorList):
 
 @pytest.mark.parametrize("dist", ['normal', 'uniform', 'sphere', 'rademacher'])
 def test_sample_like(simple_tl: TensorList, dist):
-    eps_scalar = 2.0
-    result_tl_scalar = simple_tl.sample_like(eps_scalar, distribution=dist)
+    eps_scalar = 1
+    result_tl_scalar = simple_tl.sample_like(distribution=dist)
     assert isinstance(result_tl_scalar, TensorList)
     assert result_tl_scalar.shape == simple_tl.shape
 
-    eps_list = [0.5, 1.0, 1.5]
-    result_tl_list = simple_tl.sample_like(eps_list, distribution=dist)
+    eps_list = [1.0,]
+    result_tl_list = simple_tl.sample_like(distribution=dist)
     assert isinstance(result_tl_list, TensorList)
     assert result_tl_list.shape == simple_tl.shape
 
     # Basic checks based on distribution
     if dist == 'uniform':
-        assert all(torch.all((t >= -eps_scalar/2) & (t <= eps_scalar/2)) for t in result_tl_scalar)
-        assert all(torch.all((t >= -e/2) & (t <= e/2)) for t, e in zip(result_tl_list, eps_list))
+        assert all(torch.all((t >= -eps_scalar) & (t <= eps_scalar)) for t in result_tl_scalar)
+        assert all(torch.all((t >= -e) & (t <= e)) for t, e in zip(result_tl_list, eps_list))
     elif dist == 'sphere':
-        assert torch.allclose(result_tl_scalar.global_vector_norm(), torch.tensor(eps_scalar))
+        # assert torch.allclose(result_tl_scalar.global_vector_norm(), torch.tensor(eps_scalar))
+        pass
         # Cannot check list version easily
     elif dist == 'rademacher':
          assert all(torch.all((t == -eps_scalar) | (t == eps_scalar)) for t in result_tl_scalar)

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/__init__.py

@@ -1,4 +1,4 @@
 from . import core, optim, utils
 from .core import Modular
-from .utils import compile
+from .utils import set_compilation
 from . import modules as m

torchzero/core/__init__.py

@@ -1,2 +1,2 @@
-from .module import Var, Module, Modular, Chain, maybe_chain, Chainable
-from .transform import Transform, TensorwiseTransform, Target, apply_transform
+from .module import Chain, Chainable, Modular, Module, Var, maybe_chain
+from .transform import Target, TensorwiseTransform, Transform, apply_transform