torchzero 0.3.8__tar.gz → 0.3.10__tar.gz

{torchzero-0.3.8 → torchzero-0.3.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.3.8
+Version: 0.3.10
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 License: MIT License
@@ -157,13 +157,14 @@ for epoch in range(100):
   * `NewtonCG`: Matrix-free newton's method with conjugate gradient solver.
   * `NystromSketchAndSolve`: Nyström sketch-and-solve method.
   * `NystromPCG`: NewtonCG with Nyström preconditioning (usually beats NewtonCG).
+  * `HigherOrderNewton`: Higher order Newton's method with trust region.
 
 * **Quasi-Newton**: Approximate second-order optimization methods.
   * `LBFGS`: Limited-memory BFGS.
   * `LSR1`: Limited-memory SR1.
   * `OnlineLBFGS`: Online LBFGS.
-  * `BFGS`, `SR1`, `DFP`, `BroydenGood`, `BroydenBad`, `Greenstadt1`, `Greenstadt2`, `ColumnUpdatingMethod`, `ThomasOptimalMethod`, `PSB`, `Pearson2`, `SSVM`: Classic full-matrix quasi-newton methods.
-  * `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`: Conjugate gradient methods.
+  * `BFGS`, `DFP`, `PSB`, `SR1`, `SSVM`, `BroydenBad`, `BroydenGood`, `ColumnUpdatingMethod`, `FletcherVMM`, `GradientCorrection`, `Greenstadt1`, `Greenstadt2`, `Horisho`, `McCormick`, `Pearson`, `ProjectedNewtonRaphson`, `ThomasOptimalMethod`: Classic full-matrix quasi-newton methods.
+  * `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`, `ProjectedGradientMethod`: Conjugate gradient methods.
 
 * **Line Search**:
   * `Backtracking`, `AdaptiveBacktracking`: Backtracking line searches (adaptive is my own).
@@ -312,20 +313,20 @@ not in the module itself. Also both per-parameter settings and state are stored
 
 ```python
 import torch
-from torchzero.core import Module, Vars
+from torchzero.core import Module, Var
 
 class HeavyBall(Module):
     def __init__(self, momentum: float = 0.9, dampening: float = 0):
         defaults = dict(momentum=momentum, dampening=dampening)
         super().__init__(defaults)
 
-    def step(self, vars: Vars):
-        # a module takes a Vars object, modifies it or creates a new one, and returns it
-        # Vars has a bunch of attributes, including parameters, gradients, update, closure, loss
+    def step(self, var: Var):
+        # a module takes a Var object, modifies it or creates a new one, and returns it
+        # Var has a bunch of attributes, including parameters, gradients, update, closure, loss
         # for now we are only interested in update, and we will apply the heavyball rule to it.
 
-        params = vars.params
-        update = vars.get_update() # list of tensors
+        params = var.params
+        update = var.get_update() # list of tensors
 
         exp_avg_list = []
         for p, u in zip(params, update):
@@ -346,16 +347,15 @@ class HeavyBall(Module):
             # and it is part of self.state
             exp_avg_list.append(buf.clone())
 
-        # set new update to vars
-        vars.update = exp_avg_list
-        return vars
+        # set new update to var
+        var.update = exp_avg_list
+        return var
 ```
 
 There are a some specialized base modules that make it much easier to implement some specific things.
 
 * `GradApproximator` for gradient approximations
 * `LineSearch` for line searches
-* `Preconditioner` for preconditioners
 * `Projection` for projections like GaLore or into fourier domain.
 * `QuasiNewtonH` for full-matrix quasi-newton methods that update hessian inverse approximation (because they are all very similar)
 * `ConguateGradientBase` for conjugate gradient methods, basically the only difference is how beta is calculated.
@@ -376,4 +376,4 @@ There are also wrappers providing `torch.optim.Optimizer` interface for for `sci
 
 They are in `torchzero.optim.wrappers.scipy.ScipyMinimize`, `torchzero.optim.wrappers.nlopt.NLOptOptimizer`, and `torchzero.optim.wrappers.nevergrad.NevergradOptimizer`. Make sure closure has `backward` argument as described in **Advanced Usage**.
 
-Apparently https://github.com/avaneev/biteopt is diabolical so I will add a wrapper for it too very soon.
+Apparently <https://github.com/avaneev/biteopt> is diabolical so I will add a wrapper for it too very soon.

{torchzero-0.3.8 → torchzero-0.3.10}/README.md

@@ -118,13 +118,14 @@ for epoch in range(100):
   * `NewtonCG`: Matrix-free newton's method with conjugate gradient solver.
   * `NystromSketchAndSolve`: Nyström sketch-and-solve method.
   * `NystromPCG`: NewtonCG with Nyström preconditioning (usually beats NewtonCG).
+  * `HigherOrderNewton`: Higher order Newton's method with trust region.
 
 * **Quasi-Newton**: Approximate second-order optimization methods.
   * `LBFGS`: Limited-memory BFGS.
   * `LSR1`: Limited-memory SR1.
   * `OnlineLBFGS`: Online LBFGS.
-  * `BFGS`, `SR1`, `DFP`, `BroydenGood`, `BroydenBad`, `Greenstadt1`, `Greenstadt2`, `ColumnUpdatingMethod`, `ThomasOptimalMethod`, `PSB`, `Pearson2`, `SSVM`: Classic full-matrix quasi-newton methods.
-  * `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`: Conjugate gradient methods.
+  * `BFGS`, `DFP`, `PSB`, `SR1`, `SSVM`, `BroydenBad`, `BroydenGood`, `ColumnUpdatingMethod`, `FletcherVMM`, `GradientCorrection`, `Greenstadt1`, `Greenstadt2`, `Horisho`, `McCormick`, `Pearson`, `ProjectedNewtonRaphson`, `ThomasOptimalMethod`: Classic full-matrix quasi-newton methods.
+  * `PolakRibiere`, `FletcherReeves`, `HestenesStiefel`, `DaiYuan`, `LiuStorey`, `ConjugateDescent`, `HagerZhang`, `HybridHS_DY`, `ProjectedGradientMethod`: Conjugate gradient methods.
 
 * **Line Search**:
   * `Backtracking`, `AdaptiveBacktracking`: Backtracking line searches (adaptive is my own).
@@ -273,20 +274,20 @@ not in the module itself. Also both per-parameter settings and state are stored
 
 ```python
 import torch
-from torchzero.core import Module, Vars
+from torchzero.core import Module, Var
 
 class HeavyBall(Module):
     def __init__(self, momentum: float = 0.9, dampening: float = 0):
         defaults = dict(momentum=momentum, dampening=dampening)
         super().__init__(defaults)
 
-    def step(self, vars: Vars):
-        # a module takes a Vars object, modifies it or creates a new one, and returns it
-        # Vars has a bunch of attributes, including parameters, gradients, update, closure, loss
+    def step(self, var: Var):
+        # a module takes a Var object, modifies it or creates a new one, and returns it
+        # Var has a bunch of attributes, including parameters, gradients, update, closure, loss
         # for now we are only interested in update, and we will apply the heavyball rule to it.
 
-        params = vars.params
-        update = vars.get_update() # list of tensors
+        params = var.params
+        update = var.get_update() # list of tensors
 
         exp_avg_list = []
         for p, u in zip(params, update):
@@ -307,16 +308,15 @@ class HeavyBall(Module):
             # and it is part of self.state
             exp_avg_list.append(buf.clone())
 
-        # set new update to vars
-        vars.update = exp_avg_list
-        return vars
+        # set new update to var
+        var.update = exp_avg_list
+        return var
 ```
 
 There are a some specialized base modules that make it much easier to implement some specific things.
 
 * `GradApproximator` for gradient approximations
 * `LineSearch` for line searches
-* `Preconditioner` for preconditioners
 * `Projection` for projections like GaLore or into fourier domain.
 * `QuasiNewtonH` for full-matrix quasi-newton methods that update hessian inverse approximation (because they are all very similar)
 * `ConguateGradientBase` for conjugate gradient methods, basically the only difference is how beta is calculated.
@@ -337,4 +337,4 @@ There are also wrappers providing `torch.optim.Optimizer` interface for for `sci
 
 They are in `torchzero.optim.wrappers.scipy.ScipyMinimize`, `torchzero.optim.wrappers.nlopt.NLOptOptimizer`, and `torchzero.optim.wrappers.nevergrad.NevergradOptimizer`. Make sure closure has `backward` argument as described in **Advanced Usage**.
 
-Apparently https://github.com/avaneev/biteopt is diabolical so I will add a wrapper for it too very soon.
+Apparently <https://github.com/avaneev/biteopt> is diabolical so I will add a wrapper for it too very soon.

{torchzero-0.3.8 → torchzero-0.3.10}/pyproject.toml

@@ -13,7 +13,7 @@ build-backend = "setuptools.build_meta"
 name = "torchzero"
 description = "Modular optimization library for PyTorch."
 
-version = "0.3.8"
+version = "0.3.10"
 dependencies = [
   "torch",
   "numpy",

{torchzero-0.3.8 → torchzero-0.3.10}/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()),

{torchzero-0.3.8 → torchzero-0.3.10}/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:

{torchzero-0.3.8 → torchzero-0.3.10}/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-0.3.10/torchzero/core/__init__.py

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