torchzero 0.3.14__py3-none-any.whl → 0.4.0__py3-none-any.whl

tests/test_objective.py

@@ -0,0 +1,188 @@
+import pytest
+import torch
+from torchzero.core import Objective
+from torchzero.utils.tensorlist import TensorList
+
+@torch.no_grad
+def test_get_loss():
+
+    # ---------------------------- test that it works ---------------------------- #
+    params = [torch.tensor(2.0, requires_grad=True)]
+    evaluated = False
+
+    def closure_1(backward=True):
+        assert not backward, 'backward = True'
+
+        # ensure closure only evaluates once
+        nonlocal evaluated
+        assert evaluated is False, 'closure was evaluated twice'
+        evaluated = True
+
+        loss = params[0]**2
+        if backward:
+            params[0].grad = None
+            loss.backward()
+        else:
+            assert not loss.requires_grad, "loss requires grad with backward=False"
+        return loss
+
+    obj = Objective(params=params, closure=closure_1, model=None, current_step=0)
+
+    assert obj.loss is None, obj.loss
+
+    assert (loss := obj.get_loss(backward=False)) == 4.0, loss
+    assert evaluated, evaluated
+    assert loss is obj.loss
+    assert obj.loss == 4.0
+    assert obj.loss_approx == 4.0
+    assert obj.grads is None, obj.grads
+
+    # reevaluate, which should just return already evaluated loss
+    assert (loss := obj.get_loss(backward=False)) == 4.0, loss
+    assert obj.grads is None, obj.grads
+
+
+    # ----------------------- test that backward=True works ---------------------- #
+    params = [torch.tensor(3.0, requires_grad=True)]
+    evaluated = False
+
+    def closure_2(backward=True):
+        # ensure closure only evaluates once
+        nonlocal evaluated
+        assert evaluated is False, 'closure was evaluated twice'
+        evaluated = True
+
+        loss = params[0] * 2
+        if backward:
+            assert loss.requires_grad, "loss does not require grad so `with torch.enable_grad()` context didn't work"
+            params[0].grad = None
+            loss.backward()
+        else:
+            assert not loss.requires_grad, "loss requires grad with backward=False"
+        return loss
+
+    obj = Objective(params=params, closure=closure_2, model=None, current_step=0)
+    assert obj.grads is None, obj.grads
+    assert (loss := obj.get_loss(backward=True)) == 6.0, loss
+    assert obj.grads is not None
+    assert obj.grads[0] == 2.0, obj.grads
+
+    # reevaluate, which should just return already evaluated loss
+    assert (loss := obj.get_loss(backward=True)) == 6.0, loss
+    assert obj.grads[0] == 2.0, obj.grads
+
+    # get grad, which should just return already evaluated grad
+    assert (grad := obj.get_grads())[0] == 2.0, grad
+    assert grad is obj.grads, grad
+
+    # get update, which should create and return cloned grad
+    assert obj.updates is None
+    assert (update := obj.get_updates())[0] == 2.0, update
+    assert update is obj.updates
+    assert update is not obj.grads
+    assert obj.grads is not None
+    assert update[0] == obj.grads[0]
+
+@torch.no_grad
+def test_get_grad():
+    params = [torch.tensor(2.0, requires_grad=True)]
+    evaluated = False
+
+    def closure(backward=True):
+        # ensure closure only evaluates once
+        nonlocal evaluated
+        assert evaluated is False, 'closure was evaluated twice'
+        evaluated = True
+
+        loss = params[0]**2
+        if backward:
+            assert loss.requires_grad, "loss does not require grad so `with torch.enable_grad()` context didn't work"
+            params[0].grad = None
+            loss.backward()
+        else:
+            assert not loss.requires_grad, "loss requires grad with backward=False"
+        return loss
+
+    obj = Objective(params=params, closure=closure, model=None, current_step=0)
+    assert (grad := obj.get_grads())[0] == 4.0, grad
+    assert grad is obj.grads
+
+    assert obj.loss == 4.0
+    assert (loss := obj.get_loss(backward=False)) == 4.0, loss
+    assert (loss := obj.get_loss(backward=True)) == 4.0, loss
+    assert obj.loss_approx == 4.0
+
+    assert obj.updates is None, obj.updates
+    assert (update := obj.get_updates())[0] == 4.0, update
+
+@torch.no_grad
+def test_get_update():
+    params = [torch.tensor(2.0, requires_grad=True)]
+    evaluated = False
+
+    def closure(backward=True):
+        # ensure closure only evaluates once
+        nonlocal evaluated
+        assert evaluated is False, 'closure was evaluated twice'
+        evaluated = True
+
+        loss = params[0]**2
+        if backward:
+            assert loss.requires_grad, "loss does not require grad so `with torch.enable_grad()` context didn't work"
+            params[0].grad = None
+            loss.backward()
+        else:
+            assert not loss.requires_grad, "loss requires grad with backward=False"
+        return loss
+
+    obj = Objective(params=params, closure=closure, model=None, current_step=0)
+    assert obj.updates is None, obj.updates
+    assert (update := obj.get_updates())[0] == 4.0, update
+    assert update is obj.updates
+
+    assert (grad := obj.get_grads())[0] == 4.0, grad
+    assert grad is obj.grads
+    assert grad is not update
+
+    assert obj.loss == 4.0
+    assert (loss := obj.get_loss(backward=False)) == 4.0, loss
+    assert (loss := obj.get_loss(backward=True)) == 4.0, loss
+    assert obj.loss_approx == 4.0
+
+    assert (update := obj.get_updates())[0] == 4.0, update
+
+
+def _assert_objectives_are_same_(o1: Objective, o2: Objective, clone_update: bool):
+    for k,v in o1.__dict__.items():
+        if not k.startswith('__'):
+            # if k == 'post_step_hooks': continue
+            if k == 'storage': continue
+            elif k == 'updates' and clone_update:
+                if o1.updates is None or o2.updates is None:
+                    assert o1.updates is None and o2.updates is None, f'`{k}` attribute is not the same, {o1.updates = }, {o2.updates = }'
+                else:
+                    assert (TensorList(o1.updates) == TensorList(o2.updates)).global_all()
+                    assert o1.updates is not o2.updates
+            elif k == 'params':
+                for p1, p2 in zip(o1.params, o2.params):
+                    assert p1.untyped_storage() == p2.untyped_storage()
+            else:
+                assert getattr(o2, k) is v, f'`{k}` attribute is not the same, {getattr(o1, k) = }, {getattr(o2, k) = }'
+
+def test_var_clone():
+    model = torch.nn.Sequential(torch.nn.Linear(2,2), torch.nn.Linear(2,4))
+    def closure(backward): return 1
+    obj = Objective(params=list(model.parameters()), closure=closure, model=model, current_step=0)
+
+    _assert_objectives_are_same_(obj, obj.clone(clone_updates=False), clone_update=False)
+    _assert_objectives_are_same_(obj, obj.clone(clone_updates=True), clone_update=True)
+
+    obj.grads = TensorList(torch.randn(5))
+    _assert_objectives_are_same_(obj, obj.clone(clone_updates=False), clone_update=False)
+    _assert_objectives_are_same_(obj, obj.clone(clone_updates=True), clone_update=True)
+
+    obj.updates = TensorList(torch.randn(5) * 2)
+    obj.loss = torch.randn(1)
+    obj.loss_approx = obj.loss
+    _assert_objectives_are_same_(obj, obj.clone(clone_updates=False), clone_update=False)
+    _assert_objectives_are_same_(obj, obj.clone(clone_updates=True), clone_update=True)

tests/test_opts.py

@@ -4,15 +4,23 @@ 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.
 """
+import random
 from collections.abc import Callable
 from functools import partial
 
+import numpy as np
 import pytest
 import torch
+
 import torchzero as tz
 
 PRINT = False # set to true in nbs
 
+# seed
+torch.manual_seed(0)
+np.random.seed(0)
+random.seed(0)
+
 def _booth(x, y):
     return (x + 2 * y - 7) ** 2 + (2 * x + y - 5) ** 2
 
@@ -51,7 +59,7 @@ def _run_objective(opt: tz.Modular, objective: Callable, use_closure: bool, step
     losses = []
     for i in range(steps):
         if clear and i == steps//2:
-            for m in opt.unrolled_modules: m.reset() # clear on middle step to see if there are any issues with it
+            for m in opt.flat_modules: m.reset() # clear on middle step to see if there are any issues with it
 
         if use_closure:
             def closure(backward=True):
@@ -283,8 +291,8 @@ ClipNormGrowth_additive = Run(
     sphere_steps=10, sphere_loss=10,
 )
 ClipNormGrowth_global = Run(
-    func_opt=lambda p: tz.Modular(p, tz.m.ClipNormGrowth(parameterwise=False), tz.m.LR(0.1)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.ClipNormGrowth(parameterwise=False), tz.m.LR(0.1)),
+    func_opt=lambda p: tz.Modular(p, tz.m.ClipNormGrowth(tensorwise=False), tz.m.LR(0.1)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.ClipNormGrowth(tensorwise=False), tz.m.LR(0.1)),
     needs_closure=False,
     func='booth', steps=50, loss=1e-6, merge_invariant=True,
     sphere_steps=10, sphere_loss=10,
@@ -340,56 +348,56 @@ RandomizedFDM_central2 = Run(
     sphere_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(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,
+    sphere_steps=200, sphere_loss=420,
 )
 RandomizedFDM_forward2 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='forward2', seed=0), tz.m.LR(0.01)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='forward2', 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,
+    sphere_steps=200, sphere_loss=420,
 )
 RandomizedFDM_backward2 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='backward2', seed=0), tz.m.LR(0.01)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='backward2', 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,
+    sphere_steps=200, sphere_loss=420,
 )
 RandomizedFDM_forward3 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='forward3', seed=0), tz.m.LR(0.01)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='forward3', 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,
+    sphere_steps=200, sphere_loss=420,
 )
 RandomizedFDM_backward3 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='backward3', seed=0), tz.m.LR(0.01)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='backward3', 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,
+    sphere_steps=200, sphere_loss=420,
 )
 RandomizedFDM_central4 = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='central4', seed=0), tz.m.LR(0.01)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.RandomizedFDM(formula='central4', 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,
+    sphere_steps=200, sphere_loss=420,
 )
 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,
+    sphere_steps=200, sphere_loss=420,
 )
 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,
+    sphere_steps=200, sphere_loss=420,
 )
 
 
@@ -427,35 +435,35 @@ ForwardGradient = Run(
     sphere_opt=lambda p: tz.Modular(p, tz.m.ForwardGradient(seed=0), tz.m.LR(0.001)),
     needs_closure=True,
     func='booth', steps=50, loss=40, merge_invariant=True,
-    sphere_steps=100, sphere_loss=450,
+    sphere_steps=200, sphere_loss=450,
 )
 ForwardGradient_forward = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.ForwardGradient(seed=0, jvp_method='forward'), tz.m.LR(0.01)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.ForwardGradient(seed=0, jvp_method='forward'), tz.m.LR(0.001)),
     needs_closure=True,
     func='booth', steps=50, loss=40, merge_invariant=True,
-    sphere_steps=100, sphere_loss=450,
+    sphere_steps=200, sphere_loss=450,
 )
 ForwardGradient_central = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.ForwardGradient(seed=0, jvp_method='central'), tz.m.LR(0.01)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.ForwardGradient(seed=0, jvp_method='central'), tz.m.LR(0.001)),
     needs_closure=True,
     func='booth', steps=50, loss=40, merge_invariant=True,
-    sphere_steps=100, sphere_loss=450,
+    sphere_steps=200, sphere_loss=450,
 )
 ForwardGradient_4samples = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.ForwardGradient(n_samples=4, seed=0), tz.m.LR(0.1)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.ForwardGradient(n_samples=4, seed=0), tz.m.LR(0.001)),
     needs_closure=True,
     func='booth', steps=50, loss=0.1, merge_invariant=True,
-    sphere_steps=100, sphere_loss=400,
+    sphere_steps=100, sphere_loss=420,
 )
 ForwardGradient_4samples_no_pre_generate = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.ForwardGradient(n_samples=4, seed=0, pre_generate=False), tz.m.LR(0.1)),
     sphere_opt=lambda p: tz.Modular(p, tz.m.ForwardGradient(n_samples=4, seed=0, pre_generate=False), tz.m.LR(0.001)),
     needs_closure=True,
     func='booth', steps=50, loss=0.1, merge_invariant=True,
-    sphere_steps=100, sphere_loss=400,
+    sphere_steps=100, sphere_loss=420,
 )
 
 # ------------------------- line_search/backtracking ------------------------- #
@@ -598,15 +606,15 @@ ScaleModulesByCosineSimilarity = Run(
 
 # ------------------------- momentum/matrix_momentum ------------------------- #
 MatrixMomentum_forward = Run(
-    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')),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.01, hvp_method='fd_forward'),),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='fd_forward')),
     needs_closure=True,
     func='booth', steps=50, loss=0.05, merge_invariant=True,
     sphere_steps=10, sphere_loss=0.01,
 )
 MatrixMomentum_forward = Run(
-    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')),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.01, hvp_method='fd_central')),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='fd_central')),
     needs_closure=True,
     func='booth', steps=50, loss=0.05, merge_invariant=True,
     sphere_steps=10, sphere_loss=0.01,
@@ -620,15 +628,15 @@ MatrixMomentum_forward = Run(
 )
 
 AdaptiveMatrixMomentum_forward = Run(
-    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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='fd_forward', adaptive=True)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='fd_forward', adaptive=True)),
     needs_closure=True,
     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.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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='fd_central', adaptive=True)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='fd_central', adaptive=True)),
     needs_closure=True,
     func='booth', steps=50, loss=0.05, merge_invariant=True,
     sphere_steps=10, sphere_loss=0.05,
@@ -642,15 +650,15 @@ AdaptiveMatrixMomentum_autograd = Run(
 )
 
 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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='fd_forward', adaptive=True, adapt_freq=1)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='fd_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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.05, hvp_method='fd_central', adaptive=True, adapt_freq=1)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.MatrixMomentum(0.5, hvp_method='fd_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,
@@ -720,10 +728,11 @@ Adam = Run(
 # ------------------------------ optimizers/soap ----------------------------- #
 SOAP = Run(
     func_opt=lambda p: tz.Modular(p, tz.m.SOAP(), tz.m.LR(0.4)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.SOAP(), tz.m.LR(1)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.SOAP(precond_freq=1), tz.m.LR(1)),
     needs_closure=False,
+    # merge and unmerge lrs are very different so need to test convergence separately somewhere
     func='rosen', steps=50, loss=4, merge_invariant=False,
-    sphere_steps=20, sphere_loss=25, # merge and unmerge lrs are very different so need to test convergence separately somewhere
+    sphere_steps=20, sphere_loss=25,
 )
 # ------------------------------ optimizers/lion ----------------------------- #
 Lion = Run(
@@ -735,11 +744,12 @@ 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(4)),
-    sphere_opt=lambda p: tz.Modular(p, tz.m.GraftModules(tz.m.Shampoo(), tz.m.RMSprop()), tz.m.LR(0.1)),
+    func_opt=lambda p: tz.Modular(p, tz.m.Graft(tz.m.Shampoo(), tz.m.RMSprop()), tz.m.LR(4)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.Graft(tz.m.Shampoo(), tz.m.RMSprop()), tz.m.LR(0.1)),
     needs_closure=False,
+    # 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
+    sphere_steps=20, sphere_loss=1,
 )
 
 # ------------------------- quasi_newton/quasi_newton ------------------------ #
@@ -755,6 +765,7 @@ SR1 = Run(
     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,
+    # this reaches 1e-13 on github so don't change to 0
     sphere_steps=10, sphere_loss=0,
 )
 SSVM = Run(
@@ -806,7 +817,7 @@ NewtonCG = Run(
     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,
+    func='rosen', steps=20, loss=1e-10, merge_invariant=True,
     sphere_steps=2, sphere_loss=3e-4,
 )
 
@@ -872,8 +883,8 @@ SophiaH = Run(
 
 # -------------------------- 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)),
+    func_opt=lambda p: tz.Modular(p, tz.m.experimental.HigherOrderNewton(trust_method=None)),
+    sphere_opt=lambda p: tz.Modular(p, tz.m.experimental.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,

tests/test_tensorlist.py

@@ -1567,13 +1567,6 @@ def test_where(simple_tl: TensorList):
     assert_tl_allclose(result_module, expected_tl)
 
 
-    # Test inplace where_ (needs TensorList other)
-    tl_copy = simple_tl.clone()
-    result_inplace = tl_copy.where_(condition_tl, other_tl)
-    assert result_inplace is tl_copy
-    assert_tl_allclose(tl_copy, expected_tl)
-
-
 def test_masked_fill(simple_tl: TensorList):
     mask_tl = simple_tl.lt(0)
     fill_value_scalar = 99.0
@@ -1600,7 +1593,6 @@ def test_select_set_(simple_tl: TensorList):
     mask_tl = simple_tl.gt(0.5)
     value_scalar = -1.0
     value_list_scalar = [-1.0, -2.0, -3.0]
-    value_tl = simple_tl.clone().mul_(0.1)
 
     # Set with scalar value
     tl_copy_scalar = simple_tl.clone()

tests/test_utils_optimizer.py

@@ -4,7 +4,6 @@ from functools import partial
 import pytest
 import torch
 from torchzero.utils.optimizer import (
-    Optimizer,
     get_group_vals,
     get_params,
     get_state_vals,

torchzero/__init__.py

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

torchzero/core/__init__.py

@@ -1,2 +1,8 @@
-from .module import Chain, Chainable, Modular, Module, Var, maybe_chain
-from .transform import Target, TensorwiseTransform, Transform, apply_transform
+from .transform import TensorTransform, Transform
+from .module import Chainable, Module
+from .objective import DerivativesMethod, HessianMethod, HVPMethod, Objective
+
+# order is important to avoid circular imports
+from .modular import Modular
+from .functional import apply, step, step_tensors, update
+from .chain import Chain, maybe_chain

torchzero/core/chain.py

@@ -0,0 +1,47 @@
+from collections.abc import Iterable
+
+from ..utils.python_tools import flatten
+from .module import Module, Chainable
+from .functional import _chain_step
+
+class Chain(Module):
+    """Chain modules, mostly used internally"""
+    def __init__(self, *modules: Module | Iterable[Module]):
+        super().__init__()
+        flat_modules: list[Module] = flatten(modules)
+        for i, module in enumerate(flat_modules):
+            self.set_child(f'module_{i}', module)
+
+    def update(self, objective):
+        if len(self.children) > 1:
+            raise RuntimeError("can't call `update` on Chain with more than one child, as `update` and `apply` have to be called sequentially. Use the `step` method instead of update-apply.")
+
+        if len(self.children) == 0: return
+        return self.children['module_0'].update(objective)
+
+    def apply(self, objective):
+        if len(self.children) > 1:
+            raise RuntimeError("can't call `update` on Chain with more than one child, as `update` and `apply` have to be called sequentially. Use the `step` method instead of update-apply.")
+
+        if len(self.children) == 0: return objective
+        return self.children['module_0'].apply(objective)
+
+    def step(self, objective):
+        children = [self.children[f'module_{i}'] for i in range(len(self.children))]
+        return _chain_step(objective, children)
+
+    def __repr__(self):
+        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())})'
+        return s
+
+def maybe_chain(*modules: Chainable) -> Module:
+    """Returns a single module directly if only one is provided, otherwise wraps them in a ``Chain``."""
+    flat_modules: list[Module] = flatten(modules)
+    if len(flat_modules) == 1:
+        return flat_modules[0]
+    return Chain(*flat_modules)
+
+

torchzero/core/functional.py

@@ -0,0 +1,103 @@
+from collections.abc import Mapping, Sequence, Iterable, Callable
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from .objective import Objective
+
+if TYPE_CHECKING:
+    from .module import Module
+    from .transform import Transform
+
+
+
+def update(
+    objective: "Objective",
+    module: "Transform",
+    states: list[dict[str, Any]] | None = None,
+    settings: Sequence[Mapping[str, Any]] | None = None,
+) -> None:
+    if states is None:
+        assert settings is None
+        module.update(objective)
+
+    else:
+        assert settings is not None
+        module.update_states(objective, states, settings)
+
+def apply(
+    objective: "Objective",
+    module: "Transform",
+    states: list[dict[str, Any]] | None = None,
+    settings: Sequence[Mapping[str, Any]] | None = None,
+) -> "Objective":
+    if states is None:
+        assert settings is None
+        return module.apply(objective)
+
+    else:
+        assert settings is not None
+        return module.apply_states(objective, states, settings)
+
+def _chain_step(objective: "Objective", modules: "Sequence[Module]"):
+    """steps with ``modules`` and returns updated objective, this is used within ``step`` and within ``Chain.step``"""
+    # step
+    for i, module in enumerate(modules):
+        if i!=0: objective = objective.clone(clone_updates=False)
+
+        objective = module.step(objective)
+        if objective.stop: break
+
+    return objective
+
+def step(objective: "Objective", modules: "Module | Sequence[Module]"):
+    """doesn't apply hooks!"""
+    if not isinstance(modules, Sequence):
+        modules = (modules, )
+
+    if len(modules) == 0:
+        raise RuntimeError("`modules` is an empty sequence")
+
+    # if closure is None, assume backward has been called and gather grads
+    if objective.closure is None:
+        objective.grads = [p.grad if p.grad is not None else torch.zeros_like(p) for p in objective.params]
+
+    # step and return
+    return _chain_step(objective, modules)
+
+
+def step_tensors(
+    modules: "Module | Sequence[Module]",
+    tensors: Sequence[torch.Tensor],
+    params: Iterable[torch.Tensor] | None = None,
+    grads: Sequence[torch.Tensor] | None = None,
+    loss: torch.Tensor | None = None,
+    closure: Callable | None = None,
+    objective: "Objective | None" = None
+) -> list[torch.Tensor]:
+    if objective is not None:
+        if any(i is not None for i in (params, grads, loss, closure)):
+            raise RuntimeError("Specify either `objective` or `(params, grads, loss, closure)`")
+
+    if not isinstance(modules, Sequence):
+        modules = (modules, )
+
+    # make fake params if they are only used for shapes
+    if params is None:
+        params = [t.view_as(t).requires_grad_() for t in tensors]
+
+    # create objective
+    if objective is None:
+        objective = Objective(params=params, loss=loss, closure=closure)
+
+    if grads is not None:
+        objective.grads = list(grads)
+
+    objective.updates = list(tensors)
+
+    # step with modules
+    # this won't update parameters in-place because objective.Modular is None
+    objective = _chain_step(objective, modules)
+
+    # return updates
+    return objective.get_updates()