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

torchzero/core/reformulation.py

@@ -0,0 +1,65 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+
+import torch
+
+from .module import Chainable, Modular, Module, Var
+
+
+class Reformulation(Module, ABC):
+    def __init__(self, defaults: dict | None, modules: Chainable | None):
+        super().__init__(defaults)
+
+        if modules is not None:
+            self.set_child("modules", modules)
+
+    @abstractmethod
+    def closure(self, backward: bool, closure: Callable, params:list[torch.Tensor], var: Var) -> tuple[float | torch.Tensor, Sequence[torch.Tensor] | None]:
+        """
+        returns (loss, gradient), if backward is False then gradient can be None.
+
+        If evaluating original loss/gradient at x_0, set them to ``var``.
+        """
+
+    def pre_step(self, var: Var) -> Var | None:
+        """This runs once before each step, whereas `closure` may run multiple times per step if further modules
+        evaluate gradients at multiple points. This is useful for example to pre-generate new random perturbations."""
+
+    def step(self, var):
+        ret = self.pre_step(var) # pylint:disable = assignment-from-no-return
+        if isinstance(ret, Var): var = ret
+
+        if var.closure is None: raise RuntimeError("Reformulation requires closure")
+        params, closure = var.params, var.closure
+
+        # step with children
+        if 'modules' in self.children:
+
+            # make a reformulated closure
+            def modified_closure(backward=True):
+                loss, grad = self.closure(backward, closure, params, var)
+
+                if grad is not None:
+                    for p,g in zip(params, grad):
+                        p.grad = g
+
+                return loss
+
+            # set it to a new Var object
+            modified_var = var.clone(clone_update=False)
+            modified_var.closure = modified_closure
+
+            # step with child
+            modules = self.children['modules']
+            modified_var = modules.step(modified_var)
+
+            # modified_var.loss and grad refers to loss and grad of a modified objective
+            # so we only take the update
+            var.update = modified_var.update
+
+        # or just evaluate new closure and set to update
+        else:
+            loss, grad = self.closure(backward=True, closure=closure, params=params, var=var)
+            if grad is not None: var.update = list(grad)
+
+        return var

torchzero/core/transform.py

@@ -1,18 +1,36 @@
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Sequence, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any, Literal, final
 
 import torch
 
-from ..utils import set_storage_, TensorList, vec_to_tensors
-from .module import Module, Var, Chain, Chainable
+from ..utils import TensorList, set_storage_, vec_to_tensors
+from .module import Chain, Chainable, Module, Var
 
 Target = Literal['grad', 'update', 'closure', 'params_direct', 'params_difference', 'update_difference']
 
+
 class Transform(Module, ABC):
-    """Base class for a transform. This is an abstract class, to use it, subclass it and override `update` and `apply` methods.
+    """Base class for a transform.
+    This is an abstract class, to use it, subclass it and override ``update_tensors`` and ``apply_tensors`` methods.
 
     A transform is a module that can also be applied manually to an arbitrary sequence of tensors.
+    It has two methods:
+
+    - ``update_tensors`` updates the internal state of this transform, it doesn't modify tensors. \
+            It may be called multiple times before ``apply_tensors``.
+    - ``apply_tensors`` applies this transform to tensors, without modifying the internal state if possible.
+
+    Alternatively, if update-apply structure doesn't make sense for a transform, all logic can be defined within ``apply_tensors``.
+
+    Transform can be applied to tensors corresponding to custom parameters
+    by calling ``keyed_transform_update`` and ``keyed_transform_apply``,
+    parameters will be keys to store per-parameter states, so they should remain the same python objects.
+
+    Alternatively you can manually create a list of state dictionaries per each tensor and pass it to
+    ``transform_update`` and ``transform_apply``.
+
+    A transform can modify the closure instead of directly modifying update by passing ``target="closure"``.
 
     Args:
         defaults (dict[str,Any] | None): dict with default values.
@@ -21,63 +39,63 @@ class Transform(Module, ABC):
             `grad` is always computed and can't be None. Otherwise set to False.
         target (Target, optional):
             what to set on var. Defaults to 'update'.
+
     """
     def __init__(
         self,
         defaults: dict[str,Any] | None,
-        uses_grad: bool,
+        uses_grad: bool = False,
+        uses_loss: bool = False,
         concat_params: bool = False,
         update_freq: int = 1,
-        scale_first: bool = False,
         inner: Chainable | None = None,
         target: Target = 'update',
     ):
         super().__init__(defaults)
         self._target: Target = target
         self._uses_grad = uses_grad
+        self._uses_loss = uses_loss
         self._concat_params = concat_params
         self._update_freq = update_freq
-        self._scale_first = scale_first
         self._inner = inner
+        self._var = None
 
-    def update(
+    def update_tensors(
         self,
         tensors: list[torch.Tensor],
         params: list[torch.Tensor],
         grads: list[torch.Tensor] | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         states: list[dict[str, Any]],
         settings: Sequence[Mapping[str, Any]],
     ) -> None:
-        """Updates this transform. By default does nothing - if logic is in `apply` method."""
+        """update function, this shouldn't be called directly. Updates this module."""
 
     @abstractmethod
-    def apply(
+    def apply_tensors(
         self,
         tensors: list[torch.Tensor],
         params: list[torch.Tensor],
         grads: list[torch.Tensor] | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         states: list[dict[str, Any]],
         settings: Sequence[Mapping[str, Any]],
     ) -> Sequence[torch.Tensor]:
-        """Applies the update rule to `tensors`."""
+        """apply function, this shouldn't be called directly. Applies the update rule to `tensors` and returns them.
+        If possible, this shouldn't modify the internal state of this transform."""
 
     @final
     @torch.no_grad
-    def transform(
+    def transform_update(
         self,
         tensors: list[torch.Tensor],
         params: list[torch.Tensor],
         grads: list[torch.Tensor] | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         states: list[dict[str, Any]],
         settings: Sequence[Mapping[str, Any]] | None,
-    ) -> list[torch.Tensor]:
-        """Applies this transform to an arbitrary sequence of tensors."""
-        un_tensors = tensors
-        un_params = params
-        un_grads = grads
+    ) -> None:
+        """Updates this transform from an arbitrary sequence of tensors."""
         if self._concat_params:
             tensors = [torch.cat([t.ravel() for t in tensors])]
             params = [torch.cat([p.ravel() for p in params])]
@@ -86,53 +104,67 @@ class Transform(Module, ABC):
         if settings is None:
             settings = [self.defaults for _ in tensors]
 
-        step = self.global_state.get('__step', 0)
+        step = self.global_state.get('__step', 0) # that way it gets reset correctly
+        self.global_state['__step'] = step + 1
+
         num = len(tensors)
         states = states[:num]
         settings = settings[:num]
 
-        update_freq = self._update_freq
-        scale_first = self._scale_first
-        scale_factor = 1
+        # update transform
+        if step % self._update_freq == 0:
+            self.update_tensors(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
 
-        # scaling factor for 1st step
-        if scale_first and step == 0:
-            # initial step size guess from pytorch LBFGS
-            scale_factor = 1 / TensorList(tensors).abs().global_sum().clip(min=1)
-            scale_factor = scale_factor.clip(min=torch.finfo(tensors[0].dtype).eps)
+        # store for transform_apply
+        self.global_state["__tensors"] = tensors
+        self.global_state["__params"] = params
+        self.global_state["__grads"] = grads
 
-        # update transform
-        if step % update_freq == 0:
-            self.update(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+
+    @final
+    @torch.no_grad
+    def transform_apply(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | float | None,
+        states: list[dict[str, Any]],
+        settings: Sequence[Mapping[str, Any]] | None,
+    ) -> list[torch.Tensor]:
+        """Applies this transform to an arbitrary sequence of tensors.
+        This can be used after ``transform_update`` has been used at least once."""
+
+        if settings is None:
+            settings = [self.defaults for _ in tensors]
+
+        num = len(tensors)
+        states = states[:num]
+        settings = settings[:num]
+
+        un_tensors = tensors
+        un_params = params
+        un_grads = grads
+
+        tensors = self.global_state.pop("__tensors")
+        params  = self.global_state.pop("__params")
+        grads   = self.global_state.pop("__grads")
 
         # step with inner
         if self._inner is not None:
-            tensors = apply_transform(self._inner, tensors=un_tensors, params=un_params, grads=un_grads)
+            tensors = apply_transform(self._inner, tensors=un_tensors, params=un_params, grads=un_grads, var=self._var)
             if self._concat_params:
                 tensors = [torch.cat([t.ravel() for t in tensors])]
 
         # apply transform
-        tensors = list(self.apply(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings))
-
-        # scale initial step, when preconditioner might not have been applied
-        if scale_first and step == 0:
-            torch._foreach_mul_(tensors, scale_factor)
+        tensors = list(self.apply_tensors(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings))
 
-        self.global_state['__step'] = step + 1
         if self._concat_params:
             tensors = vec_to_tensors(vec=tensors[0], reference=un_tensors)
-        return tensors
 
+        return tensors
 
-    @torch.no_grad
-    def keyed_transform(
-        self,
-        tensors: list[torch.Tensor],
-        params: list[torch.Tensor],
-        grads: list[torch.Tensor] | None,
-        loss: torch.Tensor | None,
-    ):
-        """Applies this transform to `tensors`, `params` will be used as keys and need to always point to same tensor objects."""
+    def _get_keyed_states_settings(self, params: list[torch.Tensor]):
         if self._concat_params:
             p = params[0]
             states = [self.state[p]]
@@ -145,42 +177,128 @@ class Transform(Module, ABC):
                 states.append(self.state[p])
                 settings.append(self.settings[p])
 
-        return self.transform(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+        return states, settings
+
+    @final
+    @torch.no_grad
+    def keyed_transform_update(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | float | None,
+    ):
+        """`params` will be used as keys and need to always point to same tensor objects.`"""
+        states, settings = self._get_keyed_states_settings(params)
+        self.transform_update(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+
+
+    @final
+    @torch.no_grad
+    def keyed_transform_apply(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | float | None,
+    ):
+        """`params` will be used as keys and need to always point to same tensor objects.`"""
+        states, settings = self._get_keyed_states_settings(params)
+        return self.transform_apply(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+
+
+    def pre_step(self, var: Var) -> None:
+        """Logic to run pre-transform, this way transform has access to  Var."""
+    def post_step(self, var: Var) -> None:
+        """Logic to run post-transform, this way transform has access to  Var."""
+
+    def update(self, var: Var):
+        if self._target != 'update':
+            raise ValueError("Target must be 'update' to use `update` and `apply` methods. "
+                             f"With {self._target = } only `step` method can be used.")
 
-    def step(self, var: Var) -> Var:
         # var may change, therefore current params and grads have to be extracted and passed explicitly
+        update = var.get_update() # this sets loss
         if self._uses_grad: var.get_grad()
+        if self._uses_loss: var.get_loss(False)
         params=var.params
+        self.pre_step(var)
+
+        # update
+        self._var = var
+        self.keyed_transform_update(update, params, var.grad, var.loss)
+        self._var = None
+
+    def apply(self, var: Var):
+        if self._target != 'update':
+            raise ValueError("Target must be 'update' to use `update` and `apply` methods. "
+                             f"With {self._target = } only `step` method can be used.")
+
+        # var may change, therefore current params and grads have to be extracted and passed explicitly
+        update = var.get_update() # this sets loss
+        if self._uses_grad: var.get_grad()
+        if self._uses_loss: var.get_loss(False)
+        params=var.params
+
+        # apply
+        self._var = var
+        var.update = self.keyed_transform_apply(update, params, var.grad, var.loss)
+        self._var = None
+
+        self.post_step(var)
+        return var
+
+    def step(self, var: Var) -> Var:
+
+        # var may change, therefore current params and grads have to be extracted and passed explicitly
+        if self._target in ('update', 'update_difference'): var.get_update() # this sets loss
+        if self._uses_grad or self._target == 'grad': var.get_grad()
+        if self._uses_loss: var.get_loss(False)
+        params=var.params
+        self.pre_step(var)
+        self._var = var
 
         # ---------------------------------- update ---------------------------------- #
         if self._target == 'update':
             update = var.get_update()
-            var.update = list(self.keyed_transform(update, params, var.grad, var.loss))
+            self.keyed_transform_update(update, params, var.grad, var.loss)
+            var.update = list(self.keyed_transform_apply(update, params, var.grad, var.loss))
+            self._var = None
             return var
 
         # ----------------------------------- grad ----------------------------------- #
         if self._target == 'grad':
             grad = var.get_grad()
-            var.grad = list(self.keyed_transform(grad, params, grad, var.loss))
+            self.keyed_transform_update(grad, params, grad, var.loss)
+            var.grad = list(self.keyed_transform_apply(grad, params, grad, var.loss))
+            self._var = None
             return var
 
         # ------------------------------- params_direct ------------------------------ #
         if self._target == 'params_direct':
-            new_params = self.keyed_transform(var.params, params, var.grad, var.loss)
+            self.keyed_transform_update(var.params, params, var.grad, var.loss)
+            new_params = self.keyed_transform_apply(var.params, params, var.grad, var.loss)
             for p, new_p in zip(var.params, new_params): set_storage_(p, new_p)
+            self._var = None
             return var
 
         # ----------------------------- params_differnce ----------------------------- #
         if self._target == 'params_difference':
-            new_params = tuple(self.keyed_transform([p.clone() for p in var.params], params, var.grad, var.loss))
+            p_clone = [p.clone() for p in var.params]
+            self.keyed_transform_update(p_clone, params, var.grad, var.loss)
+            new_params = tuple(self.keyed_transform_apply(p_clone, params, var.grad, var.loss))
             var.update = list(torch._foreach_sub(var.params, new_params))
+            self._var = None
             return var
 
         # ----------------------------- update_difference ---------------------------- #
         if self._target == 'update_difference':
             update = var.get_update()
-            new_update = tuple(self.keyed_transform([u.clone() for u in update], params, var.grad, var.loss))
+            u_clone = [u.clone() for u in update]
+            self.keyed_transform_update(u_clone, params, var.grad, var.loss)
+            new_update = tuple(self.keyed_transform_apply(u_clone, params, var.grad, var.loss))
             var.update = list(torch._foreach_sub(update, new_update))
+            self._var = None
             return var
 
         # ---------------------------------- closure --------------------------------- #
@@ -189,11 +307,17 @@ class Transform(Module, ABC):
             if original_closure is None: raise ValueError('Target = "closure", but closure is None')
 
             params = var.params
+            parent_var = self._var
             def transformed_closure(backward=True):
                 if backward:
                     loss = original_closure()
                     current_grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
-                    transformed_grad = list(self.keyed_transform(current_grad, params, var.grad, var.loss))
+
+                    self._var = parent_var
+                    self.keyed_transform_update(current_grad, params, var.grad, var.loss)
+                    transformed_grad = list(self.keyed_transform_apply(current_grad, params, var.grad, var.loss))
+                    self._var = None
+
                     for p, g in zip(params, transformed_grad):
                         p.grad = g
 
@@ -203,6 +327,8 @@ class Transform(Module, ABC):
                 return loss
 
             var.closure = transformed_closure
+            self.post_step(var)
+            self._var = None
             return var
 
         # ---------------------------------- invalid --------------------------------- #
@@ -212,7 +338,7 @@ class Transform(Module, ABC):
 class TensorwiseTransform(Transform, ABC):
     """Base class for a parameter-wise transform.
 
-    This is an abstract class, to use it, subclass it and override `transform`.
+    This is an abstract class, to use it, subclass it and override `update_tensor` and `apply_tensor`.
 
     Args:
         defaults (dict[str,Any] | None): dict with default values.
@@ -225,10 +351,10 @@ class TensorwiseTransform(Transform, ABC):
     def __init__(
         self,
         defaults: dict[str,Any] | None,
-        uses_grad: bool,
+        uses_grad: bool = False,
+        uses_loss: bool = False,
         concat_params: bool = False,
         update_freq: int = 1,
-        scale_first: bool = False,
         inner: Chainable | None = None,
         target: Target = 'update',
     ):
@@ -237,7 +363,7 @@ class TensorwiseTransform(Transform, ABC):
             uses_grad=uses_grad,
             concat_params=concat_params,
             update_freq=update_freq,
-            scale_first=scale_first,
+            uses_loss=uses_loss,
             inner=inner,
             target=target,
         )
@@ -247,9 +373,9 @@ class TensorwiseTransform(Transform, ABC):
         tensor: torch.Tensor,
         param: torch.Tensor,
         grad: torch.Tensor | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         state: dict[str, Any],
-        settings: Mapping[str, Any],
+        setting: Mapping[str, Any],
     ) -> None:
         """Updates this transform. By default does nothing - if logic is in `apply` method."""
 
@@ -259,20 +385,20 @@ class TensorwiseTransform(Transform, ABC):
         tensor: torch.Tensor,
         param: torch.Tensor,
         grad: torch.Tensor | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         state: dict[str, Any],
-        settings: Mapping[str, Any],
+        setting: Mapping[str, Any],
     ) -> torch.Tensor:
         """Applies the update rule to `tensor`."""
 
     @final
-    def update(self, tensors, params, grads, loss, states, settings):
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
         if grads is None: grads = [None]*len(tensors)
         for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
             self.update_tensor(t, p, g, loss, state, setting)
 
     @final
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         applied = []
         if grads is None: grads = [None]*len(tensors)
         for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
@@ -284,7 +410,7 @@ def apply_transform(
     tensors: list[torch.Tensor],
     params: list[torch.Tensor],
     grads: list[torch.Tensor] | None,
-    loss: torch.Tensor | None = None,
+    loss: torch.Tensor | float | None = None,
     var: Var | None = None,
     current_step: int = 0,
 ):
@@ -292,9 +418,10 @@ def apply_transform(
         var = Var(params=params, closure=None, model=None, current_step=current_step)
         var.loss = loss
 
-    if isinstance(tfm, Transform):
+    if isinstance(tfm, Transform) and tfm._target == 'update':
         if tfm._uses_grad and grads is None: grads = var.get_grad()
-        return list(tfm.keyed_transform(tensors, params, grads, loss))
+        tfm.keyed_transform_update(tensors, params, grads, loss)
+        return list(tfm.keyed_transform_apply(tensors, params, grads, loss))
 
     if isinstance(tfm, Chain): tfm = tfm.get_children_sequence() # pyright: ignore[reportAssignmentType]
     if isinstance(tfm, Sequence):

torchzero/modules/__init__.py

@@ -1,14 +1,23 @@
+from . import experimental
 from .clipping import *
+from .conjugate_gradient import *
 from .grad_approximation import *
+from .higher_order import *
+from .least_squares import *
 from .line_search import *
-from .lr import *
+from .misc import *
 from .momentum import *
 from .ops import *
-from .optimizers import *
+from .adaptive import *
 from .projections import *
 from .quasi_newton import *
+from .second_order import *
 from .smoothing import *
+from .step_size import *
+from .termination import *
+from .trust_region import *
+from .variance_reduction import *
 from .weight_decay import *
 from .wrappers import *
-from .second_order import *
-from .higher_order import *
+from .restarts import *
+from .zeroth_order import *

torchzero/modules/adaptive/__init__.py

@@ -0,0 +1,30 @@
+from .adagrad import Adagrad, FullMatrixAdagrad, AdagradNorm
+
+# from .curveball import CurveBall
+# from .spectral import SpectralPreconditioner
+from .adahessian import AdaHessian
+from .adam import Adam
+from .adan import Adan
+from .adaptive_heavyball import AdaptiveHeavyBall
+from .aegd import AEGD
+from .esgd import ESGD
+from .lmadagrad import LMAdagrad
+from .lion import Lion
+from .mars import MARSCorrection
+from .matrix_momentum import MatrixMomentum
+from .msam import MSAM, MSAMObjective
+from .muon import DualNormCorrection, MuonAdjustLR, Orthogonalize, orthogonalize_grads_
+from .natural_gradient import NaturalGradient
+from .orthograd import OrthoGrad, orthograd_
+from .rmsprop import RMSprop
+from .rprop import (
+    BacktrackOnSignChange,
+    Rprop,
+    ScaleLRBySignChange,
+    SignConsistencyLRs,
+    SignConsistencyMask,
+)
+from .sam import ASAM, SAM
+from .shampoo import Shampoo
+from .soap import SOAP
+from .sophia_h import SophiaH