torchzero 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

torchzero/modules/meta/grafting.py

@@ -45,15 +45,15 @@ class Graft(OptimizerModule):
 
 
     @torch.no_grad
-    def step(self, state):
-        state_copy = state.copy(clone_ascent=True)
+    def step(self, vars):
+        state_copy = vars.copy(clone_ascent=True)
         magnitude = self.children['magnitude'].return_ascent(state_copy)
 
-        if state_copy.grad is not None: state.grad = state_copy.grad
-        if state_copy.fx0 is not None: state.fx0 = state_copy.fx0
-        if state_copy.fx0_approx is not None: state.fx0_approx = state_copy.fx0_approx
+        if state_copy.grad is not None: vars.grad = state_copy.grad
+        if state_copy.fx0 is not None: vars.fx0 = state_copy.fx0
+        if state_copy.fx0_approx is not None: vars.fx0_approx = state_copy.fx0_approx
 
-        direction = self.children['direction'].return_ascent(state)
+        direction = self.children['direction'].return_ascent(vars)
 
         if self.layerwise:
             M = magnitude.norm(self.ord)
@@ -65,8 +65,8 @@ class Graft(OptimizerModule):
             D = direction.total_vector_norm(self.ord)
             if D == 0: D = M
 
-        state.ascent = direction.mul_(M / (D + self.eps))
-        return self._update_params_or_step_with_next(state)
+        vars.ascent = direction.mul_(M / (D + self.eps))
+        return self._update_params_or_step_with_next(vars)
 
 
 
@@ -94,17 +94,17 @@ class SignGrafting(OptimizerModule):
 
 
     @torch.no_grad
-    def step(self, state):
-        state_copy = state.copy(clone_ascent=True)
+    def step(self, vars):
+        state_copy = vars.copy(clone_ascent=True)
         magnitude = self.children['magnitude'].return_ascent(state_copy)
 
         # make sure to store grad and fx0 if it was calculated
-        state.update_attrs_(state_copy)
+        vars.update_attrs_(state_copy)
 
-        sign = self.children['sign'].return_ascent(state)
+        sign = self.children['sign'].return_ascent(vars)
 
-        state.ascent = magnitude.copysign_(sign)
-        return self._update_params_or_step_with_next(state)
+        vars.ascent = magnitude.copysign_(sign)
+        return self._update_params_or_step_with_next(vars)
 
 
 class IntermoduleCautious(OptimizerModule):
@@ -153,17 +153,17 @@ class IntermoduleCautious(OptimizerModule):
         self.mode: Literal["zero", "grad", "backtrack", "compare_module"]  = mode
 
     @torch.no_grad
-    def step(self, state):
+    def step(self, vars):
         params = None
-        state_copy = state.copy(clone_ascent=True)
+        state_copy = vars.copy(clone_ascent=True)
         ascent = self.children['main'].return_ascent(state_copy)
-        state.update_attrs_(state_copy)
+        vars.update_attrs_(state_copy)
 
-        if self.compare_mode == 'module': compare = self.children['compare'].return_ascent(state)
+        if self.compare_mode == 'module': compare = self.children['compare'].return_ascent(vars)
         else:
             params = self.get_params()
-            if self.compare_mode == 'ascent': compare: TensorList = state.maybe_use_grad_(params)
-            elif self.compare_mode == 'grad': compare: TensorList = state.maybe_compute_grad_(params)
+            if self.compare_mode == 'ascent': compare: TensorList = vars.maybe_use_grad_(params)
+            elif self.compare_mode == 'grad': compare: TensorList = vars.maybe_compute_grad_(params)
             else: raise ValueError(f'Invalid compare_module: {self.compare_mode}')
 
         # mask will be > 0 for parameters where both signs are the same
@@ -185,11 +185,11 @@ class IntermoduleCautious(OptimizerModule):
 
             if self.mode == 'grad':
                 params = self.get_params()
-                ascent += state.maybe_compute_grad_(params) * mask.logical_not_()
+                ascent += vars.maybe_compute_grad_(params) * mask.logical_not_()
 
             elif self.mode == 'compare_module':
                 ascent += compare * mask.logical_not_()
 
-        state.ascent = ascent
-        return self._update_params_or_step_with_next(state, params)
+        vars.ascent = ascent
+        return self._update_params_or_step_with_next(vars, params)
 

torchzero/modules/meta/optimizer_wrapper.py

@@ -64,7 +64,7 @@ class Wrap(OptimizerModule):
             self.optimizer = self._optimizer_cls(params, *self._args, **self._kwargs)
 
     @torch.no_grad
-    def step(self, state):
+    def step(self, vars):
         # check attrs
         # if self.pass_closure:
         #     if state.closure is None: raise ValueError('ClosureOptimizerWrapper requires closure.')
@@ -75,22 +75,22 @@ class Wrap(OptimizerModule):
 
         if self.next_module is None:
             # set grad to ascent and make a step with the optimizer
-            g = state.maybe_use_grad_(params)
+            g = vars.maybe_use_grad_(params)
             params.set_grad_(g)
-            state.fx0 = self.optimizer.step()
-            return state.get_loss()
+            vars.fx0 = self.optimizer.step()
+            return vars.get_loss()
 
 
         params_before_step = params.clone()
 
-        g = state.maybe_use_grad_(params)
+        g = vars.maybe_use_grad_(params)
         params.set_grad_(g)
-        state.fx0 = self.optimizer.step()
+        vars.fx0 = self.optimizer.step()
 
         # calculate update as difference in params
-        state.ascent = params_before_step - params
+        vars.ascent = params_before_step - params
         params.set_(params_before_step)
-        return self.next_module.step(state)
+        return self.next_module.step(vars)
 
 
 class WrapClosure(OptimizerModule):
@@ -148,7 +148,7 @@ class WrapClosure(OptimizerModule):
             self.optimizer = self._optimizer_cls(params, *self._args, **self._kwargs)
 
     @torch.no_grad
-    def step(self, state):
+    def step(self, vars):
         # check attrs
         # if self.pass_closure:
         #     if state.closure is None: raise ValueError('ClosureOptimizerWrapper requires closure.')
@@ -159,15 +159,15 @@ class WrapClosure(OptimizerModule):
 
         if self.next_module is None:
             # set grad to ascent and make a step with the optimizer
-            state.fx0 = self.optimizer.step(state.closure) # type:ignore
-            return state.get_loss()
+            vars.fx0 = self.optimizer.step(vars.closure) # type:ignore
+            return vars.get_loss()
 
 
         params_before_step = params.clone()
-        state.fx0 = self.optimizer.step(state.closure) # type:ignore
+        vars.fx0 = self.optimizer.step(vars.closure) # type:ignore
 
         # calculate update as difference in params
-        state.ascent = params_before_step - params
+        vars.ascent = params_before_step - params
         params.set_(params_before_step)
-        return self.next_module.step(state)
+        return self.next_module.step(vars)
 

torchzero/modules/meta/return_overrides.py

@@ -9,12 +9,12 @@ class SetGrad(OptimizerModule):
         super().__init__({})
 
     @torch.no_grad
-    def step(self, state):
+    def step(self, vars):
         if self.next_module is not None: raise ValueError("SetGrad can't have children")
         params = self.get_params()
-        g = state.maybe_use_grad_(params) # this may execute the closure which might be modified
+        g = vars.maybe_use_grad_(params) # this may execute the closure which might be modified
         params.set_grad_(g)
-        return state.get_loss()
+        return vars.get_loss()
 
 
 class ReturnAscent(OptimizerModule):
@@ -23,10 +23,10 @@ class ReturnAscent(OptimizerModule):
         super().__init__({})
 
     @torch.no_grad
-    def step(self, state) -> TensorList: # type:ignore
+    def step(self, vars) -> TensorList: # type:ignore
         if self.next_module is not None: raise ValueError("ReturnAscent can't have children")
         params = self.get_params()
-        update = state.maybe_use_grad_(params) # this will execute the closure which might be modified
+        update = vars.maybe_use_grad_(params) # this will execute the closure which might be modified
         return update
 
 class ReturnClosure(OptimizerModule):
@@ -38,9 +38,9 @@ class ReturnClosure(OptimizerModule):
         super().__init__({})
 
     @torch.no_grad
-    def step(self, state) -> _ClosureType: # type:ignore
+    def step(self, vars) -> _ClosureType: # type:ignore
         if self.next_module is not None: raise ValueError("ReturnClosure can't have children")
-        if state.closure is None:
+        if vars.closure is None:
             raise ValueError("MakeClosure requires closure")
-        return state.closure
+        return vars.closure
 

torchzero/modules/misc/accumulate.py

@@ -26,18 +26,18 @@ class Accumulate(OptimizerModule):
         self.cur_step = 0
 
     @torch.no_grad
-    def step(self, state):
+    def step(self, vars):
         self.cur_step += 1
 
         params = self.get_params()
         accumulated_update = self.get_state_key('accumulated_grads')
-        accumulated_update += state.maybe_use_grad_(params)
+        accumulated_update += vars.maybe_use_grad_(params)
 
         if self.cur_step % self.n_steps == 0:
-            state.ascent = accumulated_update.clone()
-            if self.mean: state.ascent /= self.n_steps
+            vars.ascent = accumulated_update.clone()
+            if self.mean: vars.ascent /= self.n_steps
             accumulated_update.zero_()
-            return self._update_params_or_step_with_next(state)
+            return self._update_params_or_step_with_next(vars)
 
 
-        return state.get_loss()
+        return vars.get_loss()

torchzero/modules/misc/basic.py

@@ -14,7 +14,7 @@ class Alpha(OptimizerModule):
         super().__init__(defaults)
 
     @torch.no_grad
-    def _update(self, state, ascent):
+    def _update(self, vars, ascent):
         # multiply ascent direction by lr in-place
         lr = self.get_group_key('alpha')
         ascent *= lr
@@ -27,7 +27,7 @@ class Clone(OptimizerModule):
         super().__init__({})
 
     @torch.no_grad
-    def _update(self, state, ascent): return ascent.clone()
+    def _update(self, vars, ascent): return ascent.clone()
 
 class Identity(OptimizerModule):
     """Does nothing."""
@@ -35,7 +35,7 @@ class Identity(OptimizerModule):
         super().__init__({})
 
     @torch.no_grad
-    def _update(self, state, ascent): return ascent
+    def _update(self, vars, ascent): return ascent
 
 class Lambda(OptimizerModule):
     """Applies a function to the ascent direction.
@@ -49,7 +49,7 @@ class Lambda(OptimizerModule):
         self.f = f
 
     @torch.no_grad()
-    def _update(self, state, ascent): return self.f(ascent)
+    def _update(self, vars, ascent): return self.f(ascent)
 
 class Grad(OptimizerModule):
     """Uses gradient as the update. This is useful for chains."""
@@ -57,8 +57,8 @@ class Grad(OptimizerModule):
         super().__init__({})
 
     @torch.no_grad
-    def _update(self, state, ascent):
-        ascent = state.ascent = state.maybe_compute_grad_(self.get_params())
+    def _update(self, vars, ascent):
+        ascent = vars.ascent = vars.maybe_compute_grad_(self.get_params())
         return ascent
 
 class Zeros(OptimizerModule):
@@ -66,7 +66,7 @@ class Zeros(OptimizerModule):
         super().__init__({})
 
     @torch.no_grad
-    def _update(self, state, ascent):
+    def _update(self, vars, ascent):
         return ascent.zeros_like()
 
 class Fill(OptimizerModule):
@@ -74,7 +74,7 @@ class Fill(OptimizerModule):
         super().__init__({"value": value})
 
     @torch.no_grad
-    def _update(self, state, ascent):
+    def _update(self, vars, ascent):
         return ascent.fill(self.get_group_key('value'))
 
 
@@ -83,8 +83,8 @@ class GradToUpdate(OptimizerModule):
     def __init__(self):
         super().__init__({})
 
-    def _update(self, state, ascent):
-        state.set_grad_(ascent, self.get_params())
+    def _update(self, vars, ascent):
+        vars.set_grad_(ascent, self.get_params())
         return ascent
 
 class MakeClosure(OptimizerModule):
@@ -93,12 +93,12 @@ class MakeClosure(OptimizerModule):
         super().__init__({})
         self._set_child_('modules', modules)
 
-    def step(self, state):
-        if state.closure is None: raise ValueError("MakeClosure requires a closure")
+    def step(self, vars):
+        if vars.closure is None: raise ValueError("MakeClosure requires a closure")
 
         params = self.get_params()
-        orig_closure = state.closure
-        orig_state = state.copy(True)
+        orig_closure = vars.closure
+        orig_state = vars.copy(True)
 
         def new_closure(backward = True):
             if backward:
@@ -110,6 +110,6 @@ class MakeClosure(OptimizerModule):
             else:
                 return orig_closure(False)
 
-        state.closure = new_closure # type:ignore
-        return self._update_params_or_step_with_next(state)
+        vars.closure = new_closure # type:ignore
+        return self._update_params_or_step_with_next(vars)
 

torchzero/modules/misc/lr.py

@@ -71,7 +71,7 @@ class LR(OptimizerModule):
         self._skip = False
 
     @torch.no_grad
-    def _update(self, state, ascent):
+    def _update(self, vars, ascent):
         # step with scheduler
         if self._scheduler_step_fn is not None:
             if self.cur != 0 and self.cur % self.sheduler_step_every == 0:
@@ -79,7 +79,7 @@ class LR(OptimizerModule):
 
                 # add a hook to cycle momentum
                 if self.cycle_momentum:
-                    state.add_post_step_hook(_set_momentum_hook)
+                    vars.add_post_step_hook(_set_momentum_hook)
 
             # remove init hook to delete reference to scheduler
             if self.cur == 0 and len(self.post_init_hooks) == 1:

torchzero/modules/misc/multistep.py

@@ -20,32 +20,32 @@ class Multistep(OptimizerModule):
 
         self._set_child_('modules', modules)
 
-    def step(self, state):
+    def step(self, vars):
         # no next module, just perform multiple steps
         if self.next_module is None:
             ret = None
             for step in range(self.num_steps):
-                state_copy = state.copy(clone_ascent=True) if step != self.num_steps - 1 else state
+                state_copy = vars.copy(clone_ascent=True) if step != self.num_steps - 1 else vars
                 ret = self.children['modules'].step(state_copy)
 
                 # since parameters are updated after stepping, grad and fx0 must be erased as they are no longer correct
-                state.grad = None; state.fx0 = None
+                vars.grad = None; vars.fx0 = None
 
             return ret
 
         # accumulate steps and pass to next module
         p0 = self.get_params().clone()
         for step in range(self.num_steps):
-            state_copy = state.copy(clone_ascent=True) if step != self.num_steps - 1 else state
+            state_copy = vars.copy(clone_ascent=True) if step != self.num_steps - 1 else vars
             self.children['modules'].step(state_copy)
 
             # since parameters are updated after stepping, grad and fx0 must be erased as they are no longer correct
-            state.grad = None; state.fx0 = None
+            vars.grad = None; vars.fx0 = None
 
         p1 = self.get_params()
-        state.ascent = p0 - p1
+        vars.ascent = p0 - p1
 
         # undo ascent
         p1.set_(p0)
 
-        return self._update_params_or_step_with_next(state, p1)
+        return self._update_params_or_step_with_next(vars, p1)

torchzero/modules/misc/on_increase.py

@@ -16,38 +16,38 @@ class NegateOnLossIncrease(OptimizerModule):
         self.backtrack = backtrack
 
     @torch.no_grad()
-    def step(self, state):
-        if state.closure is None: raise ValueError('NegateOnLossIncrease requires closure.')
-        if state.fx0 is None: state.fx0 = state.closure(False)
+    def step(self, vars):
+        if vars.closure is None: raise ValueError('NegateOnLossIncrease requires closure.')
+        if vars.fx0 is None: vars.fx0 = vars.closure(False)
 
         # subtract ascent direction to params and see if loss decreases
         params = self.get_params()
-        ascent_direction = state.maybe_use_grad_(params)
+        ascent_direction = vars.maybe_use_grad_(params)
         params -= ascent_direction
-        state.fx0_approx = state.closure(False)
+        vars.fx0_approx = vars.closure(False)
 
         # if this has no children, update params and return loss
         if self.next_module is None:
             if params is None: params = self.get_params()
 
-            if state.fx0_approx > state.fx0:
+            if vars.fx0_approx > vars.fx0:
                 # loss increased, so we negate thea scent direction
                 # we are currently at params - ascent direction
                 # so we add twice the ascent direction
                 params.add_(ascent_direction, alpha = 2 if self.backtrack else 1)
 
             # else: we are already at a lower loss point
-            return state.get_loss()
+            return vars.get_loss()
 
         # otherwise undo the ascent direction because it is passed to the child
         params += ascent_direction
 
         # if loss increases, negate ascent direction
-        if state.fx0_approx > state.fx0:
+        if vars.fx0_approx > vars.fx0:
             if self.backtrack: ascent_direction.neg_()
             else: ascent_direction.zero_()
 
         # otherwise undo the ascent direction and pass the updated ascent direction to the child
-        return self.next_module.step(state)
+        return self.next_module.step(vars)
 
 

torchzero/modules/momentum/momentum.py

@@ -21,7 +21,7 @@ class HeavyBall(OptimizerModule):
         super().__init__(defaults)
 
     @torch.no_grad
-    def _update(self, state, ascent):
+    def _update(self, vars, ascent):
         velocity = self.get_state_key('velocity', init = ascent)
         settings = self.get_all_group_keys()
         updated_direction = _heavyball_step(ascent, velocity, settings['momentum'], settings['dampening'])
@@ -52,7 +52,7 @@ class NesterovMomentum(OptimizerModule):
         super().__init__(defaults)
 
     @torch.no_grad
-    def _update(self, state, ascent):
+    def _update(self, vars, ascent):
         velocity = self.get_state_key('velocity')
         settings = self.get_all_group_keys()
         _nesterov_step_(ascent, velocity, settings['momentum'], settings['dampening'])
@@ -65,7 +65,7 @@ class GradientAveraging(OptimizerModule):
         super().__init__(defaults)
 
     @torch.no_grad
-    def _update(self, state, ascent):
+    def _update(self, vars, ascent):
         velocity = self.get_state_key('velocity')
         dampening = self.get_group_key('dampening')
 
@@ -89,7 +89,7 @@ class RandomCoordinateMomentum(OptimizerModule):
         self.nesterov = nesterov
 
     @torch.no_grad
-    def _update(self, state, ascent):
+    def _update(self, vars, ascent):
         velocity = self.get_state_key('velocity', init = ascent)
         settings = self.get_all_group_keys()