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

torchzero/modules/quasi_newton/sg2.py

@@ -1,6 +1,6 @@
 import torch
 
-from ...core import Module, Chainable, apply_transform
+from ...core import Module, Chainable, step
 from ...utils import TensorList, vec_to_tensors
 from ..second_order.newton import _newton_step, _get_H
 
@@ -58,12 +58,12 @@ class SG2(Module):
         if inner is not None: self.set_child('inner', inner)
 
     @torch.no_grad
-    def update(self, var):
+    def update(self, objective):
         k = self.global_state.get('step', 0) + 1
         self.global_state["step"] = k
 
-        params = TensorList(var.params)
-        closure = var.closure
+        params = TensorList(objective.params)
+        closure = objective.closure
         if closure is None:
             raise RuntimeError("closure is required for SG2")
         generator = self.get_generator(params[0].device, self.defaults["seed"])
@@ -79,7 +79,7 @@ class SG2(Module):
 
             # one sided
             if self.defaults["one_sided"]:
-                g_0 = TensorList(var.get_grad())
+                g_0 = TensorList(objective.get_grads())
                 params.add_(cd)
                 closure()
 
@@ -126,9 +126,9 @@ class SG2(Module):
 
 
     @torch.no_grad
-    def apply(self, var):
+    def apply(self, objective):
         dir = _newton_step(
-            var=var,
+            objective=objective,
             H = self.global_state["H"],
             damping = self.defaults["damping"],
             inner = self.children.get("inner", None),
@@ -138,10 +138,10 @@ class SG2(Module):
             g_proj=None,
         )
 
-        var.update = vec_to_tensors(dir, var.params)
-        return var
+        objective.updates = vec_to_tensors(dir, objective.params)
+        return objective
 
-    def get_H(self,var=...):
+    def get_H(self,objective=...):
         return _get_H(self.global_state["H"], self.defaults["eigval_fn"])
 
 
@@ -198,12 +198,12 @@ class SPSA2(Module):
         if inner is not None: self.set_child('inner', inner)
 
     @torch.no_grad
-    def update(self, var):
+    def update(self, objective):
         k = self.global_state.get('step', 0) + 1
         self.global_state["step"] = k
 
-        params = TensorList(var.params)
-        closure = var.closure
+        params = TensorList(objective.params)
+        closure = objective.closure
         if closure is None:
             raise RuntimeError("closure is required for SPSA2")
 
@@ -260,7 +260,7 @@ class SPSA2(Module):
             H_hat /= n_samples
 
         # set grad to approximated grad
-        var.grad = g_0
+        objective.grads = g_0
 
         # update H
         H = self.global_state.get("H", None)
@@ -273,9 +273,9 @@ class SPSA2(Module):
         self.global_state["H"] = H
 
     @torch.no_grad
-    def apply(self, var):
+    def apply(self, objective):
         dir = _newton_step(
-            var=var,
+            objective=objective,
             H = self.global_state["H"],
             damping = self.defaults["damping"],
             inner = self.children.get("inner", None),
@@ -285,8 +285,8 @@ class SPSA2(Module):
             g_proj=None,
         )
 
-        var.update = vec_to_tensors(dir, var.params)
-        return var
+        objective.updates = vec_to_tensors(dir, objective.params)
+        return objective
 
-    def get_H(self,var=...):
+    def get_H(self,objective=...):
         return _get_H(self.global_state["H"], self.defaults["eigval_fn"])

torchzero/modules/restarts/restars.py

@@ -4,12 +4,14 @@ from typing import final, Literal, cast
 
 import torch
 
-from ...core import Chainable, Module, Var
+from ...core import Chainable, Module, Objective
 from ...utils import TensorList
 from ..termination import TerminationCriteriaBase
 
 def _reset_except_self(optimizer, var, self: Module):
-    for m in optimizer.unrolled_modules: m.reset()
+    for m in optimizer.unrolled_modules:
+        if m is not self:
+            m.reset()
 
 class RestartStrategyBase(Module, ABC):
     """Base class for restart strategies.
@@ -24,7 +26,7 @@ class RestartStrategyBase(Module, ABC):
             self.set_child('modules', modules)
 
     @abstractmethod
-    def should_reset(self, var: Var) -> bool:
+    def should_reset(self, var: Objective) -> bool:
         """returns whether reset should occur"""
 
     def _reset_on_condition(self, var):
@@ -39,23 +41,23 @@ class RestartStrategyBase(Module, ABC):
         return modules
 
     @final
-    def update(self, var):
-        modules = self._reset_on_condition(var)
+    def update(self, objective):
+        modules = self._reset_on_condition(objective)
         if modules is not None:
-            modules.update(var)
+            modules.update(objective)
 
     @final
-    def apply(self, var):
+    def apply(self, objective):
         # don't check here because it was check in `update`
         modules = self.children.get('modules', None)
-        if modules is None: return var
-        return modules.apply(var.clone(clone_update=False))
+        if modules is None: return objective
+        return modules.apply(objective.clone(clone_updates=False))
 
     @final
-    def step(self, var):
-        modules = self._reset_on_condition(var)
-        if modules is None: return var
-        return modules.step(var.clone(clone_update=False))
+    def step(self, objective):
+        modules = self._reset_on_condition(objective)
+        if modules is None: return objective
+        return modules.step(objective.clone(clone_updates=False))
 
 
 
@@ -170,7 +172,7 @@ class PowellRestart(RestartStrategyBase):
         super().__init__(defaults, modules)
 
     def should_reset(self, var):
-        g = TensorList(var.get_grad())
+        g = TensorList(var.get_grads())
         cond1 = self.defaults['cond1']; cond2 = self.defaults['cond2']
 
         # -------------------------------- initialize -------------------------------- #
@@ -192,7 +194,7 @@ class PowellRestart(RestartStrategyBase):
 
         # ------------------------------- 2nd condition ------------------------------ #
         if (cond2 is not None) and (not reset):
-            d_g = TensorList(var.get_update()).dot(g)
+            d_g = TensorList(var.get_updates()).dot(g)
             if (-1-cond2) * g_g < d_g < (-1 + cond2) * g_g:
                 reset = True
 
@@ -229,17 +231,17 @@ class BirginMartinezRestart(Module):
 
         self.set_child("module", module)
 
-    def update(self, var):
+    def update(self, objective):
         module = self.children['module']
-        module.update(var)
+        module.update(objective)
 
-    def apply(self, var):
+    def apply(self, objective):
         module = self.children['module']
-        var = module.apply(var.clone(clone_update=False))
+        objective = module.apply(objective.clone(clone_updates=False))
 
         cond = self.defaults['cond']
-        g = TensorList(var.get_grad())
-        d = TensorList(var.get_update())
+        g = TensorList(objective.get_grads())
+        d = TensorList(objective.get_updates())
         d_g = d.dot(g)
         d_norm = d.global_vector_norm()
         g_norm = g.global_vector_norm()
@@ -247,7 +249,7 @@ class BirginMartinezRestart(Module):
         # d in our case is same direction as g so it has a minus sign
         if -d_g > -cond * d_norm * g_norm:
             module.reset()
-            var.update = g.clone()
-            return var
+            objective.updates = g.clone()
+            return objective
 
-        return var
+        return objective

torchzero/modules/second_order/__init__.py

@@ -1,7 +1,7 @@
 from .ifn import InverseFreeNewton
-from .inm import INM
+from .inm import ImprovedNewton
 from .multipoint import SixthOrder3P, SixthOrder3PM2, SixthOrder5P, TwoPointNewton
 from .newton import Newton
 from .newton_cg import NewtonCG, NewtonCGSteihaug
 from .nystrom import NystromPCG, NystromSketchAndSolve
-from .rsn import RSN
+from .rsn import SubspaceNewton

torchzero/modules/second_order/ifn.py

@@ -1,89 +1,58 @@
-import warnings
-from collections.abc import Callable
-from functools import partial
-from typing import Literal
-
 import torch
 
-from ...core import Chainable, Module, apply_transform, Var
+from ...core import Chainable, Transform, HessianMethod
 from ...utils import TensorList, vec_to_tensors
-from ...utils.linalg.linear_operator import DenseWithInverse, Dense
-from .newton import _get_H, _get_loss_grad_and_hessian, _newton_step
+from ...linalg.linear_operator import DenseWithInverse
 
 
-class InverseFreeNewton(Module):
+class InverseFreeNewton(Transform):
     """Inverse-free newton's method
 
-    .. note::
-        In most cases Newton should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Newton preconditioning to another module's output.
-
-    .. note::
-        This module requires the a closure passed to the optimizer step,
-        as it needs to re-evaluate the loss and gradients for calculating the hessian.
-        The closure must accept a ``backward`` argument (refer to documentation).
-
-    .. warning::
-        this uses roughly O(N^2) memory.
-
     Reference
         [Massalski, Marcin, and Magdalena Nockowska-Rosiak. "INVERSE-FREE NEWTON'S METHOD." Journal of Applied Analysis & Computation 15.4 (2025): 2238-2257.](https://www.jaac-online.com/article/doi/10.11948/20240428)
     """
     def __init__(
         self,
         update_freq: int = 1,
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
+        hessian_method: HessianMethod = "batched_autograd",
+        h: float = 1e-3,
         inner: Chainable | None = None,
     ):
-        defaults = dict(hessian_method=hessian_method, vectorize=vectorize, update_freq=update_freq)
-        super().__init__(defaults)
-
-        if inner is not None:
-            self.set_child('inner', inner)
+        defaults = dict(hessian_method=hessian_method, h=h)
+        super().__init__(defaults, update_freq=update_freq, inner=inner)
 
     @torch.no_grad
-    def update(self, var):
-        update_freq = self.defaults['update_freq']
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
 
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
+        _, _, H = objective.hessian(
+            hessian_method=fs['hessian_method'],
+            h=fs['h'],
+            at_x0=True
+        )
 
-        if step % update_freq == 0:
-            loss, g_list, H = _get_loss_grad_and_hessian(
-                var, self.defaults['hessian_method'], self.defaults['vectorize']
-            )
-            self.global_state["H"] = H
+        self.global_state["H"] = H
 
-            # inverse free part
-            if 'Y' not in self.global_state:
-                num = H.T
-                denom = (torch.linalg.norm(H, 1) * torch.linalg.norm(H, float('inf'))) # pylint:disable=not-callable
+        # inverse free part
+        if 'Y' not in self.global_state:
+            num = H.T
+            denom = (torch.linalg.norm(H, 1) * torch.linalg.norm(H, float('inf'))) # pylint:disable=not-callable
 
-                finfo = torch.finfo(H.dtype)
-                self.global_state['Y'] = num.div_(denom.clip(min=finfo.tiny * 2, max=finfo.max / 2))
+            finfo = torch.finfo(H.dtype)
+            self.global_state['Y'] = num.div_(denom.clip(min=finfo.tiny * 2, max=finfo.max / 2))
 
-            else:
-                Y = self.global_state['Y']
-                I2 = torch.eye(Y.size(0), device=Y.device, dtype=Y.dtype).mul_(2)
-                I2 -= H @ Y
-                self.global_state['Y'] = Y @ I2
+        else:
+            Y = self.global_state['Y']
+            I2 = torch.eye(Y.size(0), device=Y.device, dtype=Y.dtype).mul_(2)
+            I2 -= H @ Y
+            self.global_state['Y'] = Y @ I2
 
 
-    def apply(self, var):
+    def apply_states(self, objective, states, settings):
         Y = self.global_state["Y"]
-        params = var.params
-
-        # -------------------------------- inner step -------------------------------- #
-        update = var.get_update()
-        if 'inner' in self.children:
-            update = apply_transform(self.children['inner'], update, params=params, grads=var.grad, var=var)
-
-        g = torch.cat([t.ravel() for t in update])
-
-        # ----------------------------------- solve ---------------------------------- #
-        var.update = vec_to_tensors(Y@g, params)
-
-        return var
+        g = torch.cat([t.ravel() for t in objective.get_updates()])
+        objective.updates = vec_to_tensors(Y@g, objective.params)
+        return objective
 
-    def get_H(self,var):
+    def get_H(self,objective=...):
         return DenseWithInverse(A = self.global_state["H"], A_inv=self.global_state["Y"])

torchzero/modules/second_order/inm.py

@@ -1,12 +1,11 @@
 from collections.abc import Callable
-from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module
-from ...utils import TensorList, vec_to_tensors
+from ...core import Chainable, Transform, HessianMethod
+from ...utils import TensorList, vec_to_tensors, unpack_states
 from ..functional import safe_clip
-from .newton import _get_H, _get_loss_grad_and_hessian, _newton_step
+from .newton import _get_H, _newton_step
 
 @torch.no_grad
 def inm(f:torch.Tensor, J:torch.Tensor, s:torch.Tensor, y:torch.Tensor):
@@ -25,7 +24,7 @@ def _eigval_fn(J: torch.Tensor, fn) -> torch.Tensor:
     L, Q = torch.linalg.eigh(J) # pylint:disable=not-callable
     return (Q * L.unsqueeze(-2)) @ Q.mH
 
-class INM(Module):
+class ImprovedNewton(Transform):
     """Improved Newton's Method (INM).
 
     Reference:
@@ -37,69 +36,66 @@ class INM(Module):
         damping: float = 0,
         use_lstsq: bool = False,
         update_freq: int = 1,
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
-        inner: Chainable | None = None,
         H_tfm: Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, bool]] | Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
         eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        hessian_method: HessianMethod = "batched_autograd",
+        h: float = 1e-3,
+        inner: Chainable | None = None,
     ):
-        defaults = dict(damping=damping, hessian_method=hessian_method, use_lstsq=use_lstsq, vectorize=vectorize, H_tfm=H_tfm, eigval_fn=eigval_fn, update_freq=update_freq)
-        super().__init__(defaults)
-
-        if inner is not None:
-            self.set_child("inner", inner)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner'], defaults["update_freq"]
+        super().__init__(defaults, update_freq=update_freq, inner=inner, )
 
     @torch.no_grad
-    def update(self, var):
-        update_freq = self.defaults['update_freq']
-
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
-
-        if step % update_freq == 0:
-            _, f_list, J = _get_loss_grad_and_hessian(
-                var, self.defaults['hessian_method'], self.defaults['vectorize']
-            )
-
-            f = torch.cat([t.ravel() for t in f_list])
-            J = _eigval_fn(J, self.defaults["eigval_fn"])
-
-            x_list = TensorList(var.params)
-            f_list = TensorList(var.get_grad())
-            x_prev, f_prev = self.get_state(var.params, "x_prev", "f_prev", cls=TensorList)
-
-            # initialize on 1st step, do Newton step
-            if step == 0:
-                x_prev.copy_(x_list)
-                f_prev.copy_(f_list)
-                self.global_state["P"] = J
-                return
-
-            # INM update
-            s_list = x_list - x_prev
-            y_list = f_list - f_prev
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
+
+        _, f_list, J = objective.hessian(
+            hessian_method=fs['hessian_method'],
+            h=fs['h'],
+            at_x0=True
+        )
+        if f_list is None: f_list = objective.get_grads()
+
+        f = torch.cat([t.ravel() for t in f_list])
+        J = _eigval_fn(J, fs["eigval_fn"])
+
+        x_list = TensorList(objective.params)
+        f_list = TensorList(objective.get_grads())
+        x_prev, f_prev = unpack_states(states, objective.params, "x_prev", "f_prev", cls=TensorList)
+
+        # initialize on 1st step, do Newton step
+        if "P" not in self.global_state:
             x_prev.copy_(x_list)
             f_prev.copy_(f_list)
+            self.global_state["P"] = J
+            return
 
-            self.global_state["P"] = inm(f, J, s=s_list.to_vec(), y=y_list.to_vec())
+        # INM update
+        s_list = x_list - x_prev
+        y_list = f_list - f_prev
+        x_prev.copy_(x_list)
+        f_prev.copy_(f_list)
+
+        self.global_state["P"] = inm(f, J, s=s_list.to_vec(), y=y_list.to_vec())
 
 
     @torch.no_grad
-    def apply(self, var):
-        params = var.params
+    def apply_states(self, objective, states, settings):
+        fs = settings[0]
+
         update = _newton_step(
-            var=var,
+            objective = objective,
             H = self.global_state["P"],
-            damping=self.defaults["damping"],
-            inner=self.children.get("inner", None),
-            H_tfm=self.defaults["H_tfm"],
-            eigval_fn=None, # it is applied in `update`
-            use_lstsq=self.defaults["use_lstsq"],
+            damping = fs["damping"],
+            H_tfm = fs["H_tfm"],
+            eigval_fn = None, # it is applied in `update`
+            use_lstsq = fs["use_lstsq"],
         )
 
-        var.update = vec_to_tensors(update, params)
+        objective.updates = vec_to_tensors(update, objective.params)
 
-        return var
+        return objective
 
-    def get_H(self,var=...):
+    def get_H(self,objective=...):
         return _get_H(self.global_state["P"], eigval_fn=None)