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

torchzero/modules/misc/misc.py

@@ -6,7 +6,7 @@ from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, Target, TensorwiseTransform, Transform, Var
+from ...core import Chainable, Module,  TensorTransform, Transform, Objective
 from ...utils import (
     Distributions,
     Metrics,
@@ -19,15 +19,15 @@ from ...utils import (
 )
 
 
-class Previous(TensorwiseTransform):
+class Previous(TensorTransform):
     """Maintains an update from n steps back, for example if n=1, returns previous update"""
-    def __init__(self, n=1, target: Target = 'update'):
+    def __init__(self, n=1):
         defaults = dict(n=n)
-        super().__init__(uses_grad=False, defaults=defaults, target=target)
+        super().__init__(defaults=defaults)
 
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+    def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
         n = setting['n']
 
         if 'history' not in state:
@@ -38,13 +38,13 @@ class Previous(TensorwiseTransform):
         return state['history'][0]
 
 
-class LastDifference(Transform):
+class LastDifference(TensorTransform):
     """Outputs difference between past two updates."""
-    def __init__(self,target: Target = 'update'):
-        super().__init__({}, target=target)
+    def __init__(self,):
+        super().__init__()
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         prev_tensors = unpack_states(states, tensors, 'prev_tensors') # initialized to 0
         difference = torch._foreach_sub(tensors, prev_tensors)
         for p, c in zip(prev_tensors, tensors): p.set_(c)
@@ -53,16 +53,16 @@ class LastDifference(Transform):
 class LastGradDifference(Module):
     """Outputs difference between past two gradients."""
     def __init__(self):
-        super().__init__({})
+        super().__init__()
 
     @torch.no_grad
-    def step(self, var):
-        grad = var.get_grad()
-        prev_grad = self.get_state(var.params, 'prev_grad') # initialized to 0
+    def apply(self, objective):
+        grad = objective.get_grads()
+        prev_grad = self.get_state(objective.params, 'prev_grad') # initialized to 0
         difference = torch._foreach_sub(grad, prev_grad)
         for p, c in zip(prev_grad, grad): p.copy_(c)
-        var.update = list(difference)
-        return var
+        objective.updates = list(difference)
+        return objective
 
 class LastParamDifference(Module):
     """Outputs difference between past two parameters, which is the effective previous update."""
@@ -70,36 +70,36 @@ class LastParamDifference(Module):
         super().__init__({})
 
     @torch.no_grad
-    def step(self, var):
-        params = var.params
-        prev_params = self.get_state(var.params, 'prev_params') # initialized to 0
+    def apply(self, objective):
+        params = objective.params
+        prev_params = self.get_state(objective.params, 'prev_params') # initialized to 0
         difference = torch._foreach_sub(params, prev_params)
         for p, c in zip(prev_params, params): p.copy_(c)
-        var.update = list(difference)
-        return var
+        objective.updates = list(difference)
+        return objective
 
 
 
-class LastProduct(Transform):
+class LastProduct(TensorTransform):
     """Outputs difference between past two updates."""
-    def __init__(self,target: Target = 'update'):
-        super().__init__({}, uses_grad=False, target=target)
+    def __init__(self):
+        super().__init__()
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         prev = unpack_states(states, tensors, 'prev', init=torch.ones_like) # initialized to 1 for prod
         prod = torch._foreach_mul(tensors, prev)
         for p, c in zip(prev, tensors): p.set_(c)
         return prod
 
-class LastRatio(Transform):
-    """Outputs ratio between past two updates, the numerator is determined by :code:`numerator` argument."""
-    def __init__(self, numerator: Literal['cur', 'prev'] = 'cur', target: Target = 'update'):
+class LastRatio(TensorTransform):
+    """Outputs ratio between past two updates, the numerator is determined by ``numerator`` argument."""
+    def __init__(self, numerator: Literal['cur', 'prev'] = 'cur'):
         defaults = dict(numerator=numerator)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         prev = unpack_states(states, tensors, 'prev', init = torch.ones_like) # initialized to ones
         numerator = settings[0]['numerator']
         if numerator == 'cur': ratio = torch._foreach_div(tensors, prev)
@@ -107,14 +107,14 @@ class LastRatio(Transform):
         for p, c in zip(prev, tensors): p.set_(c)
         return ratio
 
-class LastAbsoluteRatio(Transform):
-    """Outputs ratio between absolute values of past two updates the numerator is determined by :code:`numerator` argument."""
-    def __init__(self, numerator: Literal['cur', 'prev'] = 'cur', eps:float=1e-8, target: Target = 'update'):
+class LastAbsoluteRatio(TensorTransform):
+    """Outputs ratio between absolute values of past two updates the numerator is determined by ``numerator`` argument."""
+    def __init__(self, numerator: Literal['cur', 'prev'] = 'cur', eps:float=1e-8):
         defaults = dict(numerator=numerator, eps=eps)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         prev = unpack_states(states, tensors, 'prev', init = torch.ones_like) # initialized to ones
         numerator = settings[0]['numerator']
         eps = NumberList(s['eps'] for s in settings)
@@ -127,139 +127,139 @@ class LastAbsoluteRatio(Transform):
         for p, c in zip(prev, tensors): p.set_(c)
         return ratio
 
-class GradSign(Transform):
+class GradSign(TensorTransform):
     """Copies gradient sign to update."""
-    def __init__(self, target: Target = 'update'):
-        super().__init__({}, uses_grad=True, target=target)
+    def __init__(self):
+        super().__init__(uses_grad=True)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         return [t.copysign_(g) for t,g in zip(tensors, grads)]
 
-class UpdateSign(Transform):
+class UpdateSign(TensorTransform):
     """Outputs gradient with sign copied from the update."""
-    def __init__(self, target: Target = 'update'):
-        super().__init__({}, uses_grad=True, target=target)
+    def __init__(self):
+        super().__init__(uses_grad=True)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         return [g.copysign(t) for t,g in zip(tensors, grads)] # no in-place
 
-class GraftToGrad(Transform):
+class GraftToGrad(TensorTransform):
     """Grafts update to the gradient, that is update is rescaled to have the same norm as the gradient."""
-    def __init__(self, tensorwise:bool=False, ord:Metrics=2, eps:float = 1e-6, target: Target = 'update'):
+    def __init__(self, tensorwise:bool=False, ord:Metrics=2, eps:float = 1e-6):
         defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps)
-        super().__init__(defaults, uses_grad=True, target=target)
+        super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(settings[0])
         return TensorList(tensors).graft_(grads, tensorwise=tensorwise, ord=ord, eps=eps)
 
-class GraftGradToUpdate(Transform):
+class GraftGradToUpdate(TensorTransform):
     """Outputs gradient grafted to update, that is gradient rescaled to have the same norm as the update."""
-    def __init__(self, tensorwise:bool=False, ord:Metrics=2, eps:float = 1e-6, target: Target = 'update'):
+    def __init__(self, tensorwise:bool=False, ord:Metrics=2, eps:float = 1e-6):
         defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps)
-        super().__init__(defaults, uses_grad=True, target=target)
+        super().__init__(defaults, uses_grad=True)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         assert grads is not None
         tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(settings[0])
         return TensorList(grads).graft(tensors, tensorwise=tensorwise, ord=ord, eps=eps)
 
 
-class GraftToParams(Transform):
-    """Grafts update to the parameters, that is update is rescaled to have the same norm as the parameters, but no smaller than :code:`eps`."""
-    def __init__(self, tensorwise:bool=False, ord:Metrics=2, eps:float = 1e-4, target: Target = 'update'):
+class GraftToParams(TensorTransform):
+    """Grafts update to the parameters, that is update is rescaled to have the same norm as the parameters, but no smaller than ``eps``."""
+    def __init__(self, tensorwise:bool=False, ord:Metrics=2, eps:float = 1e-4):
         defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(settings[0])
         return TensorList(tensors).graft_(params, tensorwise=tensorwise, ord=ord, eps=eps)
 
-class Relative(Transform):
-    """Multiplies update by absolute parameter values to make it relative to their magnitude, :code:`min_value` is minimum allowed value to avoid getting stuck at 0."""
-    def __init__(self, min_value:float = 1e-4, target: Target = 'update'):
+class Relative(TensorTransform):
+    """Multiplies update by absolute parameter values to make it relative to their magnitude, ``min_value`` is minimum allowed value to avoid getting stuck at 0."""
+    def __init__(self, min_value:float = 1e-4):
         defaults = dict(min_value=min_value)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         mul = TensorList(params).abs().clamp_([s['min_value'] for s in settings])
         torch._foreach_mul_(tensors, mul)
         return tensors
 
 class FillLoss(Module):
-    """Outputs tensors filled with loss value times :code:`alpha`"""
+    """Outputs tensors filled with loss value times ``alpha``"""
     def __init__(self, alpha: float = 1, backward: bool = True):
         defaults = dict(alpha=alpha, backward=backward)
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, var):
-        alpha = self.get_settings(var.params, 'alpha')
-        loss = var.get_loss(backward=self.defaults['backward'])
-        var.update = [torch.full_like(p, loss*a) for p,a in zip(var.params, alpha)]
-        return var
-
-class MulByLoss(Module):
-    """Multiplies update by loss times :code:`alpha`"""
-    def __init__(self, alpha: float = 1, min_value:float = 1e-8, backward: bool = True):
+    def apply(self, objective):
+        alpha = self.get_settings(objective.params, 'alpha')
+        loss = objective.get_loss(backward=self.defaults['backward'])
+        objective.updates = [torch.full_like(p, loss*a) for p,a in zip(objective.params, alpha)]
+        return objective
+
+class MulByLoss(TensorTransform):
+    """Multiplies update by loss times ``alpha``"""
+    def __init__(self, alpha: float = 1, min_value:float = 1e-16, backward: bool = True):
         defaults = dict(alpha=alpha, min_value=min_value, backward=backward)
-        super().__init__(defaults)
+        super().__init__(defaults, uses_loss=True)
 
     @torch.no_grad
-    def step(self, var):
-        alpha, min_value = self.get_settings(var.params, 'alpha', 'min_value')
-        loss = var.get_loss(backward=self.defaults['backward'])
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
+        assert loss is not None
+        alpha, min_value = unpack_dicts(settings, 'alpha', 'min_value')
         mul = [max(loss*a, mv) for a,mv in zip(alpha, min_value)]
-        torch._foreach_mul_(var.update, mul)
-        return var
+        torch._foreach_mul_(tensors, mul)
+        return tensors
 
-class DivByLoss(Module):
-    """Divides update by loss times :code:`alpha`"""
-    def __init__(self, alpha: float = 1, min_value:float = 1e-8, backward: bool = True):
+class DivByLoss(TensorTransform):
+    """Divides update by loss times ``alpha``"""
+    def __init__(self, alpha: float = 1, min_value:float = 1e-16, backward: bool = True):
         defaults = dict(alpha=alpha, min_value=min_value, backward=backward)
-        super().__init__(defaults)
+        super().__init__(defaults, uses_loss=True)
 
     @torch.no_grad
-    def step(self, var):
-        alpha, min_value = self.get_settings(var.params, 'alpha', 'min_value')
-        loss = var.get_loss(backward=self.defaults['backward'])
-        mul = [max(loss*a, mv) for a,mv in zip(alpha, min_value)]
-        torch._foreach_div_(var.update, mul)
-        return var
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
+        assert loss is not None
+        alpha, min_value = unpack_dicts(settings, 'alpha', 'min_value')
+        denom = [max(loss*a, mv) for a,mv in zip(alpha, min_value)]
+        torch._foreach_div_(tensors, denom)
+        return tensors
 
 
-class NoiseSign(Transform):
+class NoiseSign(TensorTransform):
     """Outputs random tensors with sign copied from the update."""
     def __init__(self, distribution:Distributions = 'normal', variance:float | None = None):
         defaults = dict(distribution=distribution, variance=variance)
-        super().__init__(defaults, uses_grad=False)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         variance = unpack_dicts(settings, 'variance')
         return TensorList(tensors).sample_like(settings[0]['distribution'], variance=variance).copysign_(tensors)
 
-class HpuEstimate(Transform):
+class HpuEstimate(TensorTransform):
     """returns ``y/||s||``, where ``y`` is difference between current and previous update (gradient), ``s`` is difference between current and previous parameters. The returned tensors are a finite difference approximation to hessian times previous update."""
     def __init__(self):
         defaults = dict()
-        super().__init__(defaults, uses_grad=False)
+        super().__init__(defaults)
 
     def reset_for_online(self):
         super().reset_for_online()
         self.clear_state_keys('prev_params', 'prev_update')
 
     @torch.no_grad
-    def update_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
         prev_params, prev_update = self.get_state(params, 'prev_params', 'prev_update') # initialized to 0
         s = torch._foreach_sub(params, prev_params)
         y = torch._foreach_sub(tensors, prev_update)
@@ -269,50 +269,48 @@ class HpuEstimate(Transform):
         self.store(params, 'y', y)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         return [self.state[p]['y'] for p in params]
 
 class RandomHvp(Module):
-    """Returns a hessian-vector product with a random vector"""
+    """Returns a hessian-vector product with a random vector, optionally times vector"""
 
     def __init__(
         self,
         n_samples: int = 1,
         distribution: Distributions = "normal",
         update_freq: int = 1,
-        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
+        zHz: bool = False,
+        hvp_method: Literal["autograd", "fd_forward", "central"] = "autograd",
         h=1e-3,
+        seed: int | None = None
     ):
-        defaults = dict(n_samples=n_samples, distribution=distribution, hvp_method=hvp_method, h=h, update_freq=update_freq)
+        defaults = locals().copy()
+        del defaults['self']
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        settings = self.settings[params[0]]
-        n_samples = settings['n_samples']
-        distribution = settings['distribution']
-        hvp_method = settings['hvp_method']
-        h = settings['h']
-        update_freq = settings['update_freq']
+    def apply(self, objective):
+        params = TensorList(objective.params)
 
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
         D = None
+        update_freq = self.defaults['update_freq']
         if step % update_freq == 0:
 
-            rgrad = None
-            for i in range(n_samples):
-                u = params.sample_like(distribution=distribution, variance=1)
-
-                Hvp, rgrad = var.hessian_vector_product(u, at_x0=True, rgrad=rgrad, hvp_method=hvp_method,
-                                    h=h, normalize=True, retain_graph=i < n_samples-1)
-
-                if D is None: D = Hvp
-                else: torch._foreach_add_(D, Hvp)
+            D, _ = objective.hutchinson_hessian(
+                rgrad = None,
+                at_x0 = True,
+                n_samples = self.defaults['n_samples'],
+                distribution = self.defaults['distribution'],
+                hvp_method = self.defaults['hvp_method'],
+                h = self.defaults['h'],
+                zHz = self.defaults["zHz"],
+                generator = self.get_generator(params[0].device, self.defaults["seed"]),
+            )
 
-            if n_samples > 1: torch._foreach_div_(D, n_samples)
             if update_freq != 1:
                 assert D is not None
                 D_buf = self.get_state(params, "D", cls=TensorList)
@@ -321,8 +319,8 @@ class RandomHvp(Module):
         if D is None:
             D = self.get_state(params, "D", cls=TensorList)
 
-        var.update = list(D)
-        return var
+        objective.updates = list(D)
+        return objective
 
 @torch.no_grad
 def _load_best_parameters(params: Sequence[torch.Tensor], best_params: Sequence[torch.Tensor]):
@@ -344,7 +342,7 @@ class SaveBest(Module):
         return (1 - x)**2 + (100 * (y - x**2))**2
 
     xy = torch.tensor((-1.1, 2.5), requires_grad=True)
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         [xy],
         tz.m.NAG(0.999),
         tz.m.LR(1e-6),
@@ -370,14 +368,14 @@ class SaveBest(Module):
         super().__init__()
 
     @torch.no_grad
-    def step(self, var):
-        loss = tofloat(var.get_loss(False))
+    def apply(self, objective):
+        loss = tofloat(objective.get_loss(False))
         lowest_loss = self.global_state.get('lowest_loss', float("inf"))
 
         if loss < lowest_loss:
             self.global_state['lowest_loss'] = loss
-            best_params = var.attrs['best_params'] = [p.clone() for p in var.params]
-            var.attrs['best_loss'] = loss
-            var.attrs['load_best_params'] = partial(_load_best_parameters, params=var.params, best_params=best_params)
+            best_params = objective.attrs['best_params'] = [p.clone() for p in objective.params]
+            objective.attrs['best_loss'] = loss
+            objective.attrs['load_best_params'] = partial(_load_best_parameters, params=objective.params, best_params=best_params)
 
-        return var
+        return objective