heavyball 1.2.1__tar.gz → 1.2.2__tar.gz

{heavyball-1.2.1 → heavyball-1.2.2}/PKG-INFO

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

{heavyball-1.2.1 → heavyball-1.2.2}/heavyball/chainable.py

@@ -287,16 +287,19 @@ def heavyball_momentum(group, updates, grads, params, momentum):
     return utils.heavyball_momentum(momentum, updates, utils.get_beta1(group))
 
 
+_optim_fns = {'adam': utils.adam_, 'laprop': utils.laprop_}
+
+
 @zero_guard("exp_avg", "exp_avg_sq")
 @general_guard("Q", "GG", init_fn=_init_soap)
 @no_state
-def scale_by_soap(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG):
+def scale_by_soap(group, update, grad, param, exp_avg, exp_avg_sq, Q, GG, inner: str = 'adam'):
     update = utils.promote(update)
 
     grad_projected = [utils.project(u, q, False) for u, q in zip(update, Q)]
-    precond = utils.adam_(exp_avg, exp_avg_sq, grad_projected, utils.get_beta1(group), utils.get_beta2(group),
-                          utils.scalar_guard(group['step'], exp_avg[0]))
-    precond = [utils.project(p, q, False) for p, q in zip(precond, Q)]
+    fn = _optim_fns[inner]
+    precond = fn(exp_avg, exp_avg_sq, grad_projected, utils.get_beta1(group), utils.get_beta2(group), group['step'])
+    precond = [utils.project(p, q, True) for p, q in zip(precond, Q)]
 
     for u, q, gg, eas in zip(update, Q, GG, exp_avg_sq):
         utils.update_preconditioner(u, q, gg, eas, group['max_precond_dim'], group['precondition_1d'],
@@ -355,8 +358,7 @@ def scale_by_psgd(group, update, grad, param, Q, exprs, Q_cache, cache_expr: str
     update = update.to(memory_format=torch.contiguous_format)
     Q_mat = utils.line_to_triu(Q) if group['store_triu_as_line'] else Q
     _update_psgd_precond(group, param, update, Q_mat, Q, exprs, prob)
-    out = _cached_psgd_precond_grad(cached, cache_expr, exprs, update, Q_mat, Q_cache)
-    return torch.as_strided(out, old.shape, old.stride())
+    return _cached_psgd_precond_grad(cached, cache_expr, exprs, update, Q_mat, Q_cache)
 
 
 @general_guard("Q", "exprs", ("Q_cache", None), ("cache_expr", None), init_fn=_init_psgd)
@@ -512,7 +514,7 @@ class BaseOpt(ChainOpt):
 
         fns = tuple(fns)
 
-        self.compile_step =  default(compile_step, self.compile_step)
+        self.compile_step = default(compile_step, self.compile_step)
         if default(palm, self.palm):
             fns = (palm_beta2,) + fns
         if default(gradient_clipping, self.gradient_clipping) is not None:

{heavyball-1.2.1 → heavyball-1.2.2}/heavyball.egg-info/PKG-INFO

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

{heavyball-1.2.1 → heavyball-1.2.2}/setup.py

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