heavyball 1.2.3__tar.gz → 1.3.0__tar.gz

{heavyball-1.2.3 → heavyball-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 1.2.3
+Version: 1.3.0
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler

{heavyball-1.2.3 → heavyball-1.3.0}/heavyball/chainable.py

@@ -438,14 +438,15 @@ class ChainOpt(utils.StatefulOptimizer):
 
         for param in p:
             state = self.state_(param)
-            if 'step' not in state:
-                if self.compile_step:
-                    step = utils.scalar_guard(0, param)
-                state['step'] = step
-            step = state['step'].add_(1)
+            if 'step' in state:
+                step = state['step']
+            elif self.compile_step:
+                step = utils.scalar_guard(0, param)
+            else:
+                step = 0
             break
 
-        group['step'] = step
+        group['step'] = state['step'] = step = step + 1
 
         if group['warmup_steps'] and step < group['warmup_steps']:
             group['lr'] = group['base_lr'] * step / group['warmup_steps']
@@ -494,7 +495,7 @@ class BaseOpt(ChainOpt):
     auto_fuse: bool = True
 
     def __init__(self, params, defaults, foreach: bool, gradient_clipping: str_or_fn, update_clipping: str_or_fn,
-                 palm: bool = use_default, compile_step: bool = use_default, *fns):
+                 palm: bool = use_default, *fns, compile_step: bool = use_default):
         if default(update_clipping, self.update_clipping) is None:
             if fns and self.auto_fuse:
                 args, kwargs = None, None

{heavyball-1.2.3 → heavyball-1.3.0}/heavyball/utils.py

@@ -88,7 +88,7 @@ def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1:
     except ZeroDivisionError:
         ckp1 = 0
 
-    update, parameters, z = list_guard(update, parameters, z)
+    update, parameters, z, grad = list_guard(update, parameters, z, grad)
     lr, ckp1, beta1 = scalar_guard(lr, ckp1, beta1, grad[0])
     _compilable_schedule_free_(parameters, z, ckp1, update, lr, beta1, decay, grad, caution)
     return weight_sum
@@ -912,13 +912,12 @@ def copy_stochastic_(target: Tensor, source: Tensor):
 @decorator_knowngood
 def _compilable_update_(p: List[Tensor], u: List[Tensor], decay: Tensor, lr: Tensor, caution: bool,
                         g: List[Optional[Tensor]]):
-    u = [u_.view_as(p_) for u_, p_ in zip(u, p)]
-    p32, u32 = [list(map(promote, x)) for x in [p, u]]
-
-    for p32_, u32_, g_, p_ in zip(p32, u32, g, p):  # lr is data-dependent -> can't compile a foreach
+    for u_, g_, p_ in zip(u, g, p):  # lr is data-dependent -> can't compile a foreach
+        u_ = promote(u_.view_as(p_))
+        p32_ = promote(p_)
         if caution:
-            u32_ = _compilable_cautioning(promote(g_), u32_)
-        p32_ = p32_ * (1 - decay * lr) + u32_ * -lr
+            u_ = _compilable_cautioning(promote(g_), u_)
+        p32_ = p32_ * (1 - decay * lr) + u_ * -lr
         copy_stochastic_(p_, p32_)
 
 

{heavyball-1.2.3 → heavyball-1.3.0}/heavyball.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 1.2.3
+Version: 1.3.0
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler

{heavyball-1.2.3 → heavyball-1.3.0}/setup.py

@@ -10,7 +10,7 @@ setuptools.setup(
     name='heavyball',
     license='BSD',
     description='Efficient optimizers',
-    version='1.2.3',
+    version='1.3.0',
     long_description=README,
     url='https://github.com/clashluke/heavyball',
     packages=setuptools.find_packages(),