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

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):
@@ -524,7 +530,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(
@@ -604,44 +610,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,
@@ -719,11 +725,11 @@ 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 ------------------------ #
@@ -745,7 +751,7 @@ 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()),
     needs_closure=True,
-    func='rosen', steps=50, loss=0.02, merge_invariant=True,
+    func='rosen', steps=50, loss=1e-10, merge_invariant=True,
     sphere_steps=10, sphere_loss=0,
 )
 
@@ -791,7 +797,7 @@ 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 ---------------------------- #
@@ -854,8 +860,17 @@ SophiaH = Run(
     sphere_steps=10, sphere_loss=40,
 )
 
+# -------------------------- optimizers/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,
+)
+
 # ------------------------------------ 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,7 +883,25 @@ 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,
+    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.FletcherVMM,
+    tz.m.Horisho,
+    lambda scale_first: tz.m.Horisho(scale_first=scale_first, 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()),

tests/test_tensorlist.py

@@ -835,7 +835,7 @@ def test_global_reductions(simple_tl: TensorList, global_method, vec_equiv_metho
     expected = vec_equiv_func()
 
     if isinstance(result, bool): assert result == expected
-    else: assert torch.allclose(result, expected), f"Tensors not close: {result = }, {expected = }"
+    else: assert torch.allclose(result, expected, atol=1e-4), f"Tensors not close: {result = }, {expected = }"
 
 
 def test_global_vector_norm(simple_tl: TensorList):
@@ -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:

tests/test_vars.py

@@ -1,10 +1,10 @@
 import pytest
 import torch
-from torchzero.core.module import Vars
+from torchzero.core.module import Var
 from torchzero.utils.tensorlist import TensorList
 
 @torch.no_grad
-def test_vars_get_loss():
+def test_var_get_loss():
 
     # ---------------------------- test that it works ---------------------------- #
     params = [torch.tensor(2.0, requires_grad=True)]
@@ -26,20 +26,20 @@ def test_vars_get_loss():
             assert not loss.requires_grad, "loss requires grad with backward=False"
         return loss
 
-    vars = Vars(params=params, closure=closure_1, model=None, current_step=0)
+    var = Var(params=params, closure=closure_1, model=None, current_step=0)
 
-    assert vars.loss is None, vars.loss
+    assert var.loss is None, var.loss
 
-    assert (loss := vars.get_loss(backward=False)) == 4.0, loss
+    assert (loss := var.get_loss(backward=False)) == 4.0, loss
     assert evaluated, evaluated
-    assert loss is vars.loss
-    assert vars.loss == 4.0
-    assert vars.loss_approx == 4.0
-    assert vars.grad is None, vars.grad
+    assert loss is var.loss
+    assert var.loss == 4.0
+    assert var.loss_approx == 4.0
+    assert var.grad is None, var.grad
 
     # reevaluate, which should just return already evaluated loss
-    assert (loss := vars.get_loss(backward=False)) == 4.0, loss
-    assert vars.grad is None, vars.grad
+    assert (loss := var.get_loss(backward=False)) == 4.0, loss
+    assert var.grad is None, var.grad
 
 
     # ----------------------- test that backward=True works ---------------------- #
@@ -61,30 +61,30 @@ def test_vars_get_loss():
             assert not loss.requires_grad, "loss requires grad with backward=False"
         return loss
 
-    vars = Vars(params=params, closure=closure_2, model=None, current_step=0)
-    assert vars.grad is None, vars.grad
-    assert (loss := vars.get_loss(backward=True)) == 6.0, loss
-    assert vars.grad is not None
-    assert vars.grad[0] == 2.0, vars.grad
+    var = Var(params=params, closure=closure_2, model=None, current_step=0)
+    assert var.grad is None, var.grad
+    assert (loss := var.get_loss(backward=True)) == 6.0, loss
+    assert var.grad is not None
+    assert var.grad[0] == 2.0, var.grad
 
     # reevaluate, which should just return already evaluated loss
-    assert (loss := vars.get_loss(backward=True)) == 6.0, loss
-    assert vars.grad[0] == 2.0, vars.grad
+    assert (loss := var.get_loss(backward=True)) == 6.0, loss
+    assert var.grad[0] == 2.0, var.grad
 
     # get grad, which should just return already evaluated grad
-    assert (grad := vars.get_grad())[0] == 2.0, grad
-    assert grad is vars.grad, grad
+    assert (grad := var.get_grad())[0] == 2.0, grad
+    assert grad is var.grad, grad
 
     # get update, which should create and return cloned grad
-    assert vars.update is None
-    assert (update := vars.get_update())[0] == 2.0, update
-    assert update is vars.update
-    assert update is not vars.grad
-    assert vars.grad is not None
-    assert update[0] == vars.grad[0]
+    assert var.update is None
+    assert (update := var.get_update())[0] == 2.0, update
+    assert update is var.update
+    assert update is not var.grad
+    assert var.grad is not None
+    assert update[0] == var.grad[0]
 
 @torch.no_grad
-def test_vars_get_grad():
+def test_var_get_grad():
     params = [torch.tensor(2.0, requires_grad=True)]
     evaluated = False
 
@@ -103,20 +103,20 @@ def test_vars_get_grad():
             assert not loss.requires_grad, "loss requires grad with backward=False"
         return loss
 
-    vars = Vars(params=params, closure=closure, model=None, current_step=0)
-    assert (grad := vars.get_grad())[0] == 4.0, grad
-    assert grad is vars.grad
+    var = Var(params=params, closure=closure, model=None, current_step=0)
+    assert (grad := var.get_grad())[0] == 4.0, grad
+    assert grad is var.grad
 
-    assert vars.loss == 4.0
-    assert (loss := vars.get_loss(backward=False)) == 4.0, loss
-    assert (loss := vars.get_loss(backward=True)) == 4.0, loss
-    assert vars.loss_approx == 4.0
+    assert var.loss == 4.0
+    assert (loss := var.get_loss(backward=False)) == 4.0, loss
+    assert (loss := var.get_loss(backward=True)) == 4.0, loss
+    assert var.loss_approx == 4.0
 
-    assert vars.update is None, vars.update
-    assert (update := vars.get_update())[0] == 4.0, update
+    assert var.update is None, var.update
+    assert (update := var.get_update())[0] == 4.0, update
 
 @torch.no_grad
-def test_vars_get_update():
+def test_var_get_update():
     params = [torch.tensor(2.0, requires_grad=True)]
     evaluated = False
 
@@ -135,24 +135,24 @@ def test_vars_get_update():
             assert not loss.requires_grad, "loss requires grad with backward=False"
         return loss
 
-    vars = Vars(params=params, closure=closure, model=None, current_step=0)
-    assert vars.update is None, vars.update
-    assert (update := vars.get_update())[0] == 4.0, update
-    assert update is vars.update
+    var = Var(params=params, closure=closure, model=None, current_step=0)
+    assert var.update is None, var.update
+    assert (update := var.get_update())[0] == 4.0, update
+    assert update is var.update
 
-    assert (grad := vars.get_grad())[0] == 4.0, grad
-    assert grad is vars.grad
+    assert (grad := var.get_grad())[0] == 4.0, grad
+    assert grad is var.grad
     assert grad is not update
 
-    assert vars.loss == 4.0
-    assert (loss := vars.get_loss(backward=False)) == 4.0, loss
-    assert (loss := vars.get_loss(backward=True)) == 4.0, loss
-    assert vars.loss_approx == 4.0
+    assert var.loss == 4.0
+    assert (loss := var.get_loss(backward=False)) == 4.0, loss
+    assert (loss := var.get_loss(backward=True)) == 4.0, loss
+    assert var.loss_approx == 4.0
 
-    assert (update := vars.get_update())[0] == 4.0, update
+    assert (update := var.get_update())[0] == 4.0, update
 
 
-def _assert_vars_are_same_(v1: Vars, v2: Vars, clone_update: bool):
+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
@@ -165,20 +165,20 @@ def _assert_vars_are_same_(v1: Vars, v2: Vars, clone_update: bool):
             else:
                 assert getattr(v2, k) is v, f'{k} is not the same, {v1 = }, {v2 = }'
 
-def test_vars_clone():
+def test_var_clone():
     model = torch.nn.Sequential(torch.nn.Linear(2,2), torch.nn.Linear(2,4))
     def closure(backward): return 1
-    vars = Vars(params=list(model.parameters()), closure=closure, model=model, current_step=0)
+    var = Var(params=list(model.parameters()), closure=closure, model=model, current_step=0)
 
-    _assert_vars_are_same_(vars, vars.clone(clone_update=False), clone_update=False)
-    _assert_vars_are_same_(vars, vars.clone(clone_update=True), clone_update=True)
+    _assert_var_are_same_(var, var.clone(clone_update=False), clone_update=False)
+    _assert_var_are_same_(var, var.clone(clone_update=True), clone_update=True)
 
-    vars.grad = TensorList(torch.randn(5))
-    _assert_vars_are_same_(vars, vars.clone(clone_update=False), clone_update=False)
-    _assert_vars_are_same_(vars, vars.clone(clone_update=True), clone_update=True)
+    var.grad = TensorList(torch.randn(5))
+    _assert_var_are_same_(var, var.clone(clone_update=False), clone_update=False)
+    _assert_var_are_same_(var, var.clone(clone_update=True), clone_update=True)
 
-    vars.update = TensorList(torch.randn(5) * 2)
-    vars.loss = torch.randn(1)
-    vars.loss_approx = vars.loss
-    _assert_vars_are_same_(vars, vars.clone(clone_update=False), clone_update=False)
-    _assert_vars_are_same_(vars, vars.clone(clone_update=True), clone_update=True)
+    var.update = TensorList(torch.randn(5) * 2)
+    var.loss = torch.randn(1)
+    var.loss_approx = var.loss
+    _assert_var_are_same_(var, var.clone(clone_update=False), clone_update=False)
+    _assert_var_are_same_(var, var.clone(clone_update=True), clone_update=True)

torchzero/core/__init__.py

@@ -1,3 +1,2 @@
-from .module import Vars, Module, Modular, Chain, maybe_chain, Chainable
-from .transform import Transform, TensorwiseTransform, Target, apply
-from .preconditioner import Preconditioner, TensorwisePreconditioner
+from .module import Var, Module, Modular, Chain, maybe_chain, Chainable
+from .transform import Transform, TensorwiseTransform, Target, apply_transform

torchzero/core/module.py

@@ -29,8 +29,8 @@ def _closure_backward(closure, params, retain_graph, create_graph):
         return loss
 
 # region Vars
-# ----------------------------------- vars ----------------------------------- #
-class Vars:
+# ----------------------------------- var ----------------------------------- #
+class Var:
     """
     Holds the state and context passed between optimizer modules during a step.
 
@@ -74,13 +74,13 @@ class Vars:
         """loss at a point near current point. This can be useful as some modules only calculate loss at perturbed points,
         whereas some other modules require loss strictly at current point."""
 
-        self.post_step_hooks: list[Callable[[Modular, Vars]]] = []
+        self.post_step_hooks: list[Callable[[Modular, Var]]] = []
         """list of functions to be called after optimizer step.
         The signature is:
 
         .. code:: py
 
-            def hook(optimizer: Modular, vars: Vars): ...
+            def hook(optimizer: Modular, var: Vars): ...
 
         """
 
@@ -110,7 +110,7 @@ class Vars:
         """if True, the parameters will not be updated"""
 
     def get_loss(self, backward: bool, retain_graph = None, create_graph: bool = False) -> torch.Tensor | float:
-        """Returns the loss at current parameters, computing it if it hasn't been computed already and assigning :code:`vars.loss`.
+        """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."""
 
         if self.loss is None:
@@ -143,7 +143,7 @@ class Vars:
 
     def get_grad(self, retain_graph: bool | None = None, create_graph: bool = False) -> list[torch.Tensor]:
         """Returns the gradient at initial parameters, computing it if it hasn't been computed already and assigning
-        :code:`vars.grad` and potentially :code:`vars.loss`. Do not call this at perturbed parameters."""
+        :code:`var.grad` and potentially :code:`var.loss`. Do not call this at perturbed parameters."""
         if self.grad is None:
             if self.closure is None: raise RuntimeError("closure is None")
             self.get_loss(backward=True, retain_graph=retain_graph, create_graph=create_graph) # evaluate and set self.loss and self.grad
@@ -152,15 +152,15 @@ class Vars:
         return self.grad
 
     def get_update(self) -> list[torch.Tensor]:
-        """Returns the update. If update is None, it is initialized by cloning the gradients and assigning to :code:`vars.update`.
-        Computing the gradients may assign :code:`vars.grad` and :code:`vars.loss` if they haven't been computed.
+        """Returns the update. If update is None, it is initialized by cloning the gradients and assigning to :code:`var.update`.
+        Computing the gradients may assign :code:`var.grad` and :code:`var.loss` if they haven't been computed.
         Do not call this at perturbed parameters."""
         if self.update is None: self.update = [g.clone() for g in self.get_grad()]
         return self.update
 
     def clone(self, clone_update: bool):
         """Creates a shallow copy of the Vars object, update can optionally be deep-copied (via :code:`torch.clone`)."""
-        copy = Vars(params = self.params, closure=self.closure, model=self.model, current_step=self.current_step)
+        copy = Var(params = self.params, closure=self.closure, model=self.model, current_step=self.current_step)
 
         if clone_update and self.update is not None:
             copy.update = [u.clone() for u in self.update]
@@ -176,16 +176,16 @@ class Vars:
 
         return copy
 
-    def update_attrs_from_clone_(self, vars: "Vars"):
+    def update_attrs_from_clone_(self, var: "Var"):
         """Updates attributes of this `Vars` instance from a cloned instance.
         Typically called after a child module has processed a cloned `Vars`
         object. This propagates any newly computed loss or gradient values
         from the child's context back to the parent `Vars` if the parent
         didn't have them computed already.
         """
-        if self.loss is None: self.loss = vars.loss
-        if self.loss_approx is None: self.loss_approx = vars.loss_approx
-        if self.grad is None: self.grad = vars.grad
+        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
 
     def zero_grad(self, set_to_none=True):
         if set_to_none:
@@ -269,36 +269,36 @@ class Module(ABC):
         return s
 
     @overload
-    def get_settings(self, key: str, *,
-                     params: Sequence[torch.Tensor], cls: type[ListLike] = list) -> ListLike: ...
+    def get_settings(self, params: Sequence[torch.Tensor], key: str, *,
+                     cls: type[ListLike] = list) -> ListLike: ...
     @overload
-    def get_settings(self, key: list[str] | tuple[str,...], *,
-                     params: Sequence[torch.Tensor], cls: type[ListLike] = list) -> list[ListLike]: ...
+    def get_settings(self, params: Sequence[torch.Tensor], key: list[str] | tuple[str,...], *,
+                     cls: type[ListLike] = list) -> list[ListLike]: ...
     @overload
-    def get_settings(self, key: str, key2: str, *keys: str,
-                     params: Sequence[torch.Tensor], cls: type[ListLike] = list) -> list[ListLike]: ...
+    def get_settings(self, params: Sequence[torch.Tensor], key: str, key2: str, *keys: str,
+                     cls: type[ListLike] = list) -> list[ListLike]: ...
 
-    def get_settings(self, key: str | list[str] | tuple[str,...], key2: str | None = None, *keys: str,
-                     params: Sequence[torch.Tensor], cls: type[ListLike] = list) -> ListLike | list[ListLike]:
+    def get_settings(self, params: Sequence[torch.Tensor], key: str | list[str] | tuple[str,...], key2: str | None = None,
+                     *keys: str, cls: type[ListLike] = list) -> ListLike | list[ListLike]:
         # if isinstance(params, Vars): params = params.params
         return get_state_vals(self.settings, params, key, key2, *keys, must_exist=True, cls=cls) # pyright:ignore[reportArgumentType]
 
 
     @overload
-    def get_state(self, key: str, *,
-                   params: Sequence[torch.Tensor], must_exist: bool = False, init: Init = torch.zeros_like,
+    def get_state(self, params: Sequence[torch.Tensor], key: str, *,
+                   must_exist: bool = False, init: Init = torch.zeros_like,
                    cls: type[ListLike] = list) -> ListLike: ...
     @overload
-    def get_state(self, key: list[str] | tuple[str,...], *,
-                   params: Sequence[torch.Tensor], must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+    def get_state(self, params: Sequence[torch.Tensor], key: list[str] | tuple[str,...], *,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
                    cls: type[ListLike] = list) -> list[ListLike]: ...
     @overload
-    def get_state(self, key: str, key2: str, *keys: str,
-                   params: Sequence[torch.Tensor], must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+    def get_state(self, params: Sequence[torch.Tensor], key: str, key2: str, *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
                    cls: type[ListLike] = list) -> list[ListLike]: ...
 
-    def get_state(self, key: str | list[str] | tuple[str,...], key2: str | None = None, *keys: str,
-                   params: Sequence[torch.Tensor], must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
+    def get_state(self, params: Sequence[torch.Tensor], key: str | list[str] | tuple[str,...], key2: str | None = None, *keys: str,
+                   must_exist: bool = False, init: Init | Sequence[Init] = torch.zeros_like,
                    cls: type[ListLike] = list) -> ListLike | list[ListLike]:
         """Returns values of per-parameter state for a given key.
         If key doesn't exist, create it with inits.
@@ -404,8 +404,8 @@ class Module(ABC):
 
     # ---------------------------- OVERRIDABLE METHODS --------------------------- #
     @abstractmethod
-    def step(self, vars: Vars) -> Vars:
-        """performs a step, returns new vars but may update them in-place."""
+    def step(self, var: Var) -> Var:
+        """performs a step, returns new var but may update them in-place."""
 
     def reset(self):
         """Resets the internal state of the module (e.g. momentum)."""
@@ -556,13 +556,13 @@ class Modular(torch.optim.Optimizer):
                 if not p.requires_grad: continue
                 for map in self._per_parameter_global_settings[p]: map.update(settings)
 
-        # create vars
+        # create var
         params = [p for g in self.param_groups for p in g['params'] if p.requires_grad]
-        vars = Vars(params=params, closure=closure, model=self.model, current_step=self.current_step)
+        var = Var(params=params, closure=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:
-            vars.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+            var.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
 
         last_module = self.modules[-1]
         last_lr = last_module.defaults.get('lr', None)
@@ -570,27 +570,27 @@ class Modular(torch.optim.Optimizer):
 
         # step
         for i, module in enumerate(self.modules):
-            if i!=0: vars = vars.clone(clone_update=False)
+            if i!=0: var = var.clone(clone_update=False)
 
             # last module, or next to last module before lr
             if (i == n_modules - 1) or ((i == n_modules - 2) and (last_lr is not None)):
-                if module.children: vars.nested_is_last = True
-                else: vars.is_last = True
-                if last_lr is not None: vars.last_module_lrs = last_module.get_settings('lr', params=vars.params)
+                if module.children: var.nested_is_last = True
+                else: var.is_last = True
+                if last_lr is not None: var.last_module_lrs = [last_module.settings[p]['lr'] for p in var.params]
 
-            vars = module.step(vars)
-            if vars.stop: break
+            var = module.step(var)
+            if var.stop: break
 
         # apply update
-        if not vars.skip_update:
+        if not var.skip_update:
             with torch.no_grad():
-                torch._foreach_sub_(params, vars.get_update())
+                torch._foreach_sub_(params, var.get_update())
 
-        for hook in vars.post_step_hooks:
-            hook(self, vars)
+        for hook in var.post_step_hooks:
+            hook(self, var)
 
         self.current_step += 1
-        return vars.loss if vars.loss is not None else vars.loss_approx
+        return var.loss if var.loss is not None else var.loss_approx
 
     def __repr__(self):
         return f'Modular({", ".join(str(m) for m in self.modules)})'
@@ -606,11 +606,11 @@ class Chain(Module):
         for i, module in enumerate(flat_modules):
             self.set_child(f'module_{i}', module)
 
-    def step(self, vars):
+    def step(self, var):
         for i in range(len(self.children)):
-            vars = self.children[f'module_{i}'].step(vars)
-            if vars.stop: break
-        return vars
+            var = self.children[f'module_{i}'].step(var)
+            if var.stop: break
+        return var
 
     def __repr__(self):
         s = self.__class__.__name__