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

torchzero/modules/optimizers/sophia_h.py

@@ -2,7 +2,7 @@ from typing import Literal
 from collections.abc import Callable
 import torch
 
-from ...core import Module, Target, Transform, Chainable, apply
+from ...core import Module, Target, Transform, Chainable, apply_transform
 from ...utils import NumberList, TensorList, as_tensorlist
 from ...utils.derivatives import hvp, hvp_fd_forward, hvp_fd_central
 
@@ -56,8 +56,8 @@ class SophiaH(Module):
             self.set_child('inner', inner)
 
     @torch.no_grad
-    def step(self, vars):
-        params = vars.params
+    def step(self, var):
+        params = var.params
         settings = self.settings[params[0]]
         hvp_method = settings['hvp_method']
         fd_h = settings['fd_h']
@@ -71,15 +71,15 @@ class SophiaH(Module):
                 self.global_state['generator'] = torch.Generator(params[0].device).manual_seed(seed)
             generator = self.global_state['generator']
 
-        beta1, beta2, precond_scale, clip, eps = self.get_settings(
-            'beta1', 'beta2', 'precond_scale', 'clip', 'eps', params=params, cls=NumberList)
+        beta1, beta2, precond_scale, clip, eps = self.get_settings(params,
+            'beta1', 'beta2', 'precond_scale', 'clip', 'eps', cls=NumberList)
 
-        exp_avg, h_exp_avg = self.get_state('exp_avg', 'h_exp_avg', params=params, cls=TensorList)
+        exp_avg, h_exp_avg = self.get_state(params, 'exp_avg', 'h_exp_avg', cls=TensorList)
 
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
-        closure = vars.closure
+        closure = var.closure
         assert closure is not None
 
         h = None
@@ -90,12 +90,12 @@ class SophiaH(Module):
                 u = [torch.randn(p.shape, device=p.device, dtype=p.dtype, generator=generator) for p in params]
 
                 if hvp_method == 'autograd':
-                    if grad is None: grad = vars.get_grad(create_graph=True)
+                    if grad is None: grad = var.get_grad(create_graph=True)
                     assert grad is not None
                     Hvp = hvp(params, grad, u, retain_graph=i < n_samples-1)
 
                 elif hvp_method == 'forward':
-                    loss, Hvp = hvp_fd_forward(closure, params, u, h=fd_h, g_0=vars.get_grad(), normalize=True)
+                    loss, Hvp = hvp_fd_forward(closure, params, u, h=fd_h, g_0=var.get_grad(), normalize=True)
 
                 elif hvp_method == 'central':
                     loss, Hvp = hvp_fd_central(closure, params, u, h=fd_h, normalize=True)
@@ -109,11 +109,11 @@ class SophiaH(Module):
             assert h is not None
             if n_samples > 1: torch._foreach_div_(h, n_samples)
 
-        update = vars.get_update()
+        update = var.get_update()
         if 'inner' in self.children:
-            update = apply(self.children['inner'], tensors=update, params=params, grads=vars.grad, vars=vars)
+            update = apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var)
 
-        vars.update = sophia_H(
+        var.update = sophia_H(
             tensors=TensorList(update),
             h=TensorList(h) if h is not None else None,
             exp_avg_=exp_avg,
@@ -126,4 +126,4 @@ class SophiaH(Module):
             eps=eps,
             step=step,
         )
-        return vars
+        return var

torchzero/modules/projections/dct.py

@@ -34,8 +34,8 @@ class DCTProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad, defaults=defaults)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
-        settings = self.settings[vars.params[0]]
+    def project(self, tensors, var, current):
+        settings = self.settings[var.params[0]]
         dims = settings['dims']
         norm = settings['norm']
 
@@ -54,8 +54,8 @@ class DCTProjection(Projection):
         return projected
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
-        settings = self.settings[vars.params[0]]
+    def unproject(self, tensors, var, current):
+        settings = self.settings[var.params[0]]
         dims = settings['dims']
         norm = settings['norm']
 

torchzero/modules/projections/fft.py

@@ -45,8 +45,8 @@ class FFTProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad, defaults=defaults)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
-        settings = self.settings[vars.params[0]]
+    def project(self, tensors, var, current):
+        settings = self.settings[var.params[0]]
         one_d = settings['one_d']
         norm = settings['norm']
 
@@ -60,14 +60,14 @@ class FFTProjection(Projection):
         return [torch.view_as_real(torch.fft.rfftn(t, norm=norm)) if t.numel() > 1 else t for t in tensors] # pylint:disable=not-callable
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
-        settings = self.settings[vars.params[0]]
+    def unproject(self, tensors, var, current):
+        settings = self.settings[var.params[0]]
         one_d = settings['one_d']
         norm = settings['norm']
 
         if one_d:
             vec = torch.view_as_complex(tensors[0])
             unprojected_vec = torch.fft.irfft(vec, n=self.global_state['length'], norm=norm) # pylint:disable=not-callable
-            return vec_to_tensors(unprojected_vec, reference=vars.params)
+            return vec_to_tensors(unprojected_vec, reference=var.params)
 
-        return [torch.fft.irfftn(torch.view_as_complex(t.contiguous()), s=p.shape, norm=norm) if t.numel() > 1 else t for t, p in zip(tensors, vars.params)] # pylint:disable=not-callable
+        return [torch.fft.irfftn(torch.view_as_complex(t.contiguous()), s=p.shape, norm=norm) if t.numel() > 1 else t for t, p in zip(tensors, var.params)] # pylint:disable=not-callable

torchzero/modules/projections/galore.py

@@ -6,5 +6,5 @@ from typing import Any, Literal
 
 import torch
 
-from ...core import Chainable, Module, Vars
+from ...core import Chainable, Module, Var
 from .projection import Projection

torchzero/modules/projections/projection.py

@@ -6,15 +6,15 @@ from typing import Any, Literal
 import warnings
 import torch
 
-from ...core import Chainable, Module, Vars
+from ...core import Chainable, Module, Var
 from ...utils import vec_to_tensors
 
 
-def _make_projected_closure(closure, vars: Vars, projection: "Projection",
+def _make_projected_closure(closure, var: Var, projection: "Projection",
                            params: list[torch.Tensor], projected_params: list[torch.Tensor]):
 
     def projected_closure(backward=True):
-        unprojected_params = projection.unproject(projected_params, vars, current='params')
+        unprojected_params = projection.unproject(projected_params, var, current='params')
 
         with torch.no_grad():
             for p, new_p in zip(params, unprojected_params):
@@ -23,7 +23,7 @@ def _make_projected_closure(closure, vars: Vars, projection: "Projection",
         if backward:
             loss = closure()
             grads = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
-            projected_grads = projection.project(grads, vars, current='grads')
+            projected_grads = projection.project(grads, var, current='grads')
             for p, g in zip(projected_params, projected_grads):
                 p.grad = g
 
@@ -38,15 +38,15 @@ def _projected_get_grad_override(
     retain_graph: bool | None = None,
     create_graph: bool = False,
     projection: Any = ...,
-    unprojected_vars: Any = ...,
+    unprojected_var: Any = ...,
     self: Any = ...,
 ):
     assert isinstance(projection, Projection)
-    assert isinstance(unprojected_vars, Vars)
-    assert isinstance(self, Vars)
+    assert isinstance(unprojected_var, Var)
+    assert isinstance(self, Var)
 
     if self.grad is not None: return self.grad
-    grads = unprojected_vars.get_grad(retain_graph, create_graph)
+    grads = unprojected_var.get_grad(retain_graph, create_graph)
     projected_grads = list(projection.project(grads, self, current='grads'))
     self.grad = projected_grads
     for p, g in zip(self.params, projected_grads):
@@ -85,56 +85,56 @@ class Projection(Module, ABC):
         self._projected_params = None
 
     @abstractmethod
-    def project(self, tensors: list[torch.Tensor], vars: Vars, current: Literal['params', 'grads', 'update']) -> Iterable[torch.Tensor]:
+    def project(self, tensors: list[torch.Tensor], var: Var, current: Literal['params', 'grads', 'update']) -> Iterable[torch.Tensor]:
         """projects `tensors`. Note that this can be called multiple times per step with `params`, `grads`, and `update`."""
 
     @abstractmethod
-    def unproject(self, tensors: list[torch.Tensor], vars: Vars, current: Literal['params', 'grads', 'update']) -> Iterable[torch.Tensor]:
+    def unproject(self, tensors: list[torch.Tensor], var: Var, current: Literal['params', 'grads', 'update']) -> Iterable[torch.Tensor]:
         """unprojects `tensors`. Note that this can be called multiple times per step with `params`, `grads`, and `update`."""
 
     @torch.no_grad
-    def step(self, vars: Vars):
-        projected_vars = vars.clone(clone_update=False)
+    def step(self, var: Var):
+        projected_var = var.clone(clone_update=False)
         update_is_grad = False
 
         # closure will calculate projected update and grad if needed
-        if self._project_params and vars.closure is not None:
-            if self._project_update and vars.update is not None: projected_vars.update = list(self.project(vars.update, vars=vars, current='update'))
+        if self._project_params and var.closure is not None:
+            if self._project_update and var.update is not None: projected_var.update = list(self.project(var.update, var=var, current='update'))
             else:
                 update_is_grad = True
-            if self._project_grad and vars.grad is not None: projected_vars.grad = list(self.project(vars.grad, vars=vars, current='grads'))
+            if self._project_grad and var.grad is not None: projected_var.grad = list(self.project(var.grad, var=var, current='grads'))
 
         # project update and grad, unprojected attributes are deleted
         else:
             if self._project_update:
-                if vars.update is None:
+                if var.update is None:
                     # update is None, meaning it will be set to `grad`.
                     # we can project grad and use it for update
-                    grad = vars.get_grad()
-                    projected_vars.grad = list(self.project(grad, vars=vars, current='grads'))
-                    if self._project_grad: projected_vars.update = [g.clone() for g in projected_vars.grad]
-                    else: projected_vars.update = projected_vars.grad.copy() # don't clone because grad shouldn't be used
-                    del vars.update
+                    grad = var.get_grad()
+                    projected_var.grad = list(self.project(grad, var=var, current='grads'))
+                    if self._project_grad: projected_var.update = [g.clone() for g in projected_var.grad]
+                    else: projected_var.update = projected_var.grad.copy() # don't clone because grad shouldn't be used
+                    del var.update
                     update_is_grad = True
 
                 else:
-                    update = vars.get_update()
-                    projected_vars.update = list(self.project(update, vars=vars, current='update'))
-                    del update, vars.update
+                    update = var.get_update()
+                    projected_var.update = list(self.project(update, var=var, current='update'))
+                    del update, var.update
 
-            if self._project_grad and projected_vars.grad is None:
-                grad = vars.get_grad()
-                projected_vars.grad = list(self.project(grad, vars=vars, current='grads'))
+            if self._project_grad and projected_var.grad is None:
+                grad = var.get_grad()
+                projected_var.grad = list(self.project(grad, var=var, current='grads'))
 
         original_params = None
         if self._project_params:
-            original_params = [p.clone() for p in vars.params]
-            projected_params = self.project(vars.params, vars=vars, current='params')
+            original_params = [p.clone() for p in var.params]
+            projected_params = self.project(var.params, var=var, current='params')
 
         else:
             # make fake params for correct shapes and state storage
             # they reuse update or grad storage for memory efficiency
-            projected_params = projected_vars.update if projected_vars.update is not None else projected_vars.grad
+            projected_params = projected_var.update if projected_var.update is not None else projected_var.grad
             assert projected_params is not None
 
         if self._projected_params is None:
@@ -148,22 +148,22 @@ class Projection(Module, ABC):
 
         # project closure
         if self._project_params:
-            closure = vars.closure; params = vars.params
-            projected_vars.closure = _make_projected_closure(closure, vars=vars, projection=self, params=params,
+            closure = var.closure; params = var.params
+            projected_var.closure = _make_projected_closure(closure, var=var, projection=self, params=params,
                                                              projected_params=self._projected_params)
 
         else:
-            projected_vars.closure = None
+            projected_var.closure = None
 
         # step
-        projected_vars.params = self._projected_params
-        projected_vars.get_grad = partial(
+        projected_var.params = self._projected_params
+        projected_var.get_grad = partial(
             _projected_get_grad_override,
             projection=self,
-            unprojected_vars=vars,
-            self=projected_vars,
+            unprojected_var=var,
+            self=projected_var,
         )
-        projected_vars = self.children['modules'].step(projected_vars)
+        projected_var = self.children['modules'].step(projected_var)
 
         # empty fake params storage
         # this doesn't affect update/grad because it is a different python object, set_ changes storage on an object
@@ -172,28 +172,28 @@ class Projection(Module, ABC):
                 p.set_(torch.empty(0, device=p.device, dtype=p.dtype)) # pyright: ignore[reportArgumentType]
 
         # unproject
-        unprojected_vars = projected_vars.clone(clone_update=False)
-        unprojected_vars.closure = vars.closure
-        unprojected_vars.params = vars.params
-        unprojected_vars.grad = vars.grad
+        unprojected_var = projected_var.clone(clone_update=False)
+        unprojected_var.closure = var.closure
+        unprojected_var.params = var.params
+        unprojected_var.grad = var.grad
 
         if self._project_update:
-            assert projected_vars.update is not None
-            unprojected_vars.update = list(self.unproject(projected_vars.update, vars=vars, current='grads' if update_is_grad else 'update'))
-            del projected_vars.update
+            assert projected_var.update is not None
+            unprojected_var.update = list(self.unproject(projected_var.update, var=var, current='grads' if update_is_grad else 'update'))
+            del projected_var.update
 
         # unprojecting grad doesn't make sense?
         # if self._project_grad:
-        #     assert projected_vars.grad is not None
-        #     unprojected_vars.grad = list(self.unproject(projected_vars.grad, vars=vars))
+        #     assert projected_var.grad is not None
+        #     unprojected_var.grad = list(self.unproject(projected_var.grad, var=var))
 
-        del projected_vars
+        del projected_var
 
         if original_params is not None:
-            for p, o in zip(unprojected_vars.params, original_params):
+            for p, o in zip(unprojected_var.params, original_params):
                 p.set_(o) # pyright: ignore[reportArgumentType]
 
-        return unprojected_vars
+        return unprojected_var
 
 
 
@@ -206,12 +206,12 @@ class FlipConcatProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
+    def project(self, tensors, var, current):
         return [torch.cat([u.view(-1) for u in tensors], dim=-1).flip(0)]
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
-        return vec_to_tensors(vec=tensors[0].flip(0), reference=vars.params)
+    def unproject(self, tensors, var, current):
+        return vec_to_tensors(vec=tensors[0].flip(0), reference=var.params)
 
 
 class NoopProjection(Projection):
@@ -221,11 +221,11 @@ class NoopProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
+    def project(self, tensors, var, current):
         return tensors
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
+    def unproject(self, tensors, var, current):
         return tensors
 
 class MultipyProjection(Projection):
@@ -235,10 +235,10 @@ class MultipyProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
+    def project(self, tensors, var, current):
         return torch._foreach_mul(tensors, 2)
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
+    def unproject(self, tensors, var, current):
         return torch._foreach_div(tensors, 2)
 

torchzero/modules/projections/structural.py

@@ -17,12 +17,12 @@ class VectorProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
+    def project(self, tensors, var, current):
         return [torch.cat([u.view(-1) for u in tensors], dim=-1)]
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
-        return vec_to_tensors(vec=tensors[0], reference=vars.params)
+    def unproject(self, tensors, var, current):
+        return vec_to_tensors(vec=tensors[0], reference=var.params)
 
 
 
@@ -33,8 +33,8 @@ class TensorizeProjection(Projection):
         super().__init__(modules, defaults=defaults, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
-        params = vars.params
+    def project(self, tensors, var, current):
+        params = var.params
         max_side = self.settings[params[0]]['max_side']
         num_elems = sum(t.numel() for t in tensors)
 
@@ -60,12 +60,12 @@ class TensorizeProjection(Projection):
         return [vec.view(dims)]
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
+    def unproject(self, tensors, var, current):
         remainder = self.global_state['remainder']
         # warnings.warn(f'{tensors[0].shape = }')
         vec = tensors[0].view(-1)
         if remainder > 0: vec = vec[:-remainder]
-        return vec_to_tensors(vec, vars.params)
+        return vec_to_tensors(vec, var.params)
 
 class BlockPartition(Projection):
     """splits parameters into blocks (for now flatttens them and chunks)"""
@@ -74,9 +74,9 @@ class BlockPartition(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad, defaults=defaults)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
+    def project(self, tensors, var, current):
         partitioned = []
-        for p,t in zip(vars.params, tensors):
+        for p,t in zip(var.params, tensors):
             settings = self.settings[p]
             max_size = settings['max_size']
             n = t.numel()
@@ -101,10 +101,10 @@ class BlockPartition(Projection):
         return partitioned
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
+    def unproject(self, tensors, var, current):
         ti = iter(tensors)
         unprojected = []
-        for p in vars.params:
+        for p in var.params:
             settings = self.settings[p]
             n = p.numel()
 
@@ -130,19 +130,19 @@ class TensorNormsProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
     @torch.no_grad
-    def project(self, tensors, vars, current):
-        orig = self.get_state(f'{current}_orig', params=vars.params)
+    def project(self, tensors, var, current):
+        orig = self.get_state(var.params, f'{current}_orig')
         torch._foreach_copy_(orig, tensors)
 
         norms = torch._foreach_norm(tensors)
-        self.get_state(f'{current}_orig_norms', params=vars.params, init=norms, cls=TensorList).set_(norms)
+        self.get_state(var.params, f'{current}_orig_norms', cls=TensorList).set_(norms)
 
         return [torch.stack(norms)]
 
     @torch.no_grad
-    def unproject(self, tensors, vars, current):
-        orig = self.get_state(f'{current}_orig', params=vars.params)
-        orig_norms = torch.stack(self.get_state(f'{current}_orig_norms', params=vars.params))
+    def unproject(self, tensors, var, current):
+        orig = self.get_state(var.params, f'{current}_orig')
+        orig_norms = torch.stack(self.get_state(var.params, f'{current}_orig_norms'))
         target_norms = tensors[0]
 
         orig_norms = torch.where(orig_norms == 0, 1, orig_norms)

torchzero/modules/quasi_newton/__init__.py

@@ -1,7 +1,36 @@
-from .cg import PolakRibiere, FletcherReeves, HestenesStiefel, DaiYuan, LiuStorey, ConjugateDescent, HagerZhang, HybridHS_DY
+from .cg import (
+    ConjugateDescent,
+    DaiYuan,
+    FletcherReeves,
+    HagerZhang,
+    HestenesStiefel,
+    HybridHS_DY,
+    LiuStorey,
+    PolakRibiere,
+    ProjectedGradientMethod,
+)
 from .lbfgs import LBFGS
+from .lsr1 import LSR1
 from .olbfgs import OnlineLBFGS
-# from .experimental import ModularLBFGS
 
-from .quasi_newton import BFGS, SR1, DFP, BroydenGood, BroydenBad, Greenstadt1, Greenstadt2, ColumnUpdatingMethod, ThomasOptimalMethod, PSB, Pearson2, SSVM
-from .lsr1 import LSR1
+# from .experimental import ModularLBFGS
+from .quasi_newton import (
+    BFGS,
+    DFP,
+    PSB,
+    SR1,
+    SSVM,
+    BroydenBad,
+    BroydenGood,
+    ColumnUpdatingMethod,
+    FletcherVMM,
+    GradientCorrection,
+    Greenstadt1,
+    Greenstadt2,
+    Horisho,
+    McCormick,
+    NewSSM,
+    Pearson,
+    ProjectedNewtonRaphson,
+    ThomasOptimalMethod,
+)