heavyball 0.17.2__py3-none-any.whl → 0.18.0__py3-none-any.whl

heavyball/cached_delayed_psgd_kron.py

@@ -42,7 +42,7 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32'):
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -50,12 +50,8 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
         if not 0.0 <= weight_decay:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
 
-        if preconditioner_update_probability is None:
-            preconditioner_update_probability = precond_update_prob_schedule()
         if clip_fn is None:
             clip_fn = lambda x: trust_region_clip_(x, 0.9, 1.5)
-        self.preconditioner_update_probability = preconditioner_update_probability
-        self.clip_fn = clip_fn
 
         defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
@@ -63,20 +59,11 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
                         # precond lr hardcoded to 0.1
                         precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
                         step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
-                        store_triu_as_line=store_triu_as_line,
-                        q_dtype=q_dtype)
-        super().__init__(params, defaults, foreach)
+                        store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
+        super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
-        self._prob_step = 0
 
     def _step(self, group):
-        # update preconditioners all together
-        update_prob = self.preconditioner_update_probability
-        if callable(update_prob):
-            update_prob = update_prob(self._prob_step)
-        do_update = self.rng.random() < update_prob
-        self._prob_step += 1
-
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
@@ -128,11 +115,11 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
             q_orig = Q_list.pop(0)
             ea = exp_avg_list.pop(0)
 
-            if do_update:
+            if self.should_update(group):
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig
                 q32 = [promote(q_) for q_ in q]
-                self.balance([g], [q32])
-                self.do_update([p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig], store_triu_as_line=store_triu_as_line)
+                self.do_update(group, [p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig],
+                               store_triu_as_line)
                 for c_, q_ in zip(cached_q, q):
                     if q_.ndim == 2:
                         torch.matmul(q_.T.conj(), q_, out=c_)

heavyball/cached_psgd_kron.py

@@ -40,7 +40,7 @@ class ForeachCachedPSGDKron(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32'):
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -48,12 +48,8 @@ class ForeachCachedPSGDKron(PSGDBase):
         if not 0.0 <= weight_decay:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
 
-        if preconditioner_update_probability is None:
-            preconditioner_update_probability = precond_update_prob_schedule()
         if clip_fn is None:
             clip_fn = lambda x: trust_region_clip_(x, 0.9, 1.5)
-        self.preconditioner_update_probability = preconditioner_update_probability
-        self.clip_fn = clip_fn
 
         defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
@@ -63,18 +59,10 @@ class ForeachCachedPSGDKron(PSGDBase):
                         step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
                         store_triu_as_line=store_triu_as_line,
                         q_dtype=q_dtype)
-        super().__init__(params, defaults, foreach)
+        super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
-        self._prob_step = 0
 
     def _step(self, group):
-        # update preconditioners all together
-        update_prob = self.preconditioner_update_probability
-        if callable(update_prob):
-            update_prob = update_prob(self._prob_step)
-        do_update = self.rng.random() < update_prob
-        self._prob_step += 1
-
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
@@ -128,11 +116,10 @@ class ForeachCachedPSGDKron(PSGDBase):
 
             new = torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea)
 
-            if do_update:
+            if self.should_update(group):
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig
                 q32 = [promote(q_) for q_ in q]
-                self.balance([g], [q32])
-                self.do_update([p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig], store_triu_as_line=store_triu_as_line)
+                self.do_update(group,[p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig], store_triu_as_line)
                 for c_, q_ in zip(cached_q, q):
                     if q_.ndim == 2:
                         torch.matmul(q_.T.conj(), q_, out=c_)

heavyball/delayed_psgd.py

@@ -39,7 +39,7 @@ class ForeachDelayedPSGD(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: callable = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32'):
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -47,12 +47,8 @@ class ForeachDelayedPSGD(PSGDBase):
         if not 0.0 <= weight_decay:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
 
-        if preconditioner_update_probability is None:
-            preconditioner_update_probability = precond_update_prob_schedule()
         if clip_fn is None:
             clip_fn = lambda x: trust_region_clip_(x, 0.9, 1.5)
-        self.preconditioner_update_probability = preconditioner_update_probability
-        self.clip_fn = clip_fn
 
         defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
@@ -61,18 +57,10 @@ class ForeachDelayedPSGD(PSGDBase):
                         precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
                         step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
                         store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
-        super().__init__(params, defaults, foreach)
+        super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
-        self._prob_step = 0
 
     def _step(self, group):
-        # update preconditioners all together
-        update_prob = self.preconditioner_update_probability
-        if callable(update_prob):
-            update_prob = update_prob(self._prob_step)
-        do_update = self.rng.random() < update_prob
-        self._prob_step += 1
-
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
@@ -114,10 +102,9 @@ class ForeachDelayedPSGD(PSGDBase):
             ea = exp_avg_list.pop(0)
             q = line_to_triu(q_orig) if store_triu_as_line else q_orig
             new = psgd_precond_grad(q, self.state_(p)["exprs"], ea)
-            if do_update:
+            if self.should_update(group):
                 q32 = [promote(q_) for q_ in q]
-                self.do_update([p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig], store_triu_as_line=store_triu_as_line)
-                self.balance([g], [q32])
+                self.do_update(group,[p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig], store_triu_as_line)
             set_(g, new)
 
         grad_list = self.clip_fn(grad_list)

heavyball/p_adam.py

@@ -5,7 +5,7 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 """
 
 import torch
-from heavyball.utils import triu_to_line, line_to_triu
+from heavyball.utils import triu_to_line, line_to_triu, identity
 
 from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, PSGDBase, precond_update_prob_schedule, \
     exp_avg_sq_, beta_debias, split_p_and_g_in_group, promote
@@ -38,8 +38,8 @@ class ForeachPaLMPAdam(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, betas=(None, None), beta: float = 0.9,
                  beta2_scale: float = 0.8, merge_dims: bool = False, split: bool = False, clip_fn: callable = None,
-                 store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32'):
+                 store_triu_as_line: bool = True, foreach: bool = True, q_dtype='float32',
+                 stochastic_schedule: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= weight_decay:
@@ -47,12 +47,8 @@ class ForeachPaLMPAdam(PSGDBase):
         if betas[0] is not None:
             beta = betas[0]
 
-        if preconditioner_update_probability is None:
-            preconditioner_update_probability = precond_update_prob_schedule()
         if clip_fn is None:
-            clip_fn = lambda x: x
-        self.preconditioner_update_probability = preconditioner_update_probability
-        self.clip_fn = clip_fn
+            clip_fn = identity
 
         defaults = dict(lr=lr, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
@@ -61,18 +57,9 @@ class ForeachPaLMPAdam(PSGDBase):
                         precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
                         step=0, warmup_steps=warmup_steps, beta=beta, beta2_scale=beta2_scale, merge_dims=merge_dims,
                         split=split, store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
-        super().__init__(params, defaults, foreach)
-
-        self._prob_step = 0
+        super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
-        # update preconditioners all together
-        update_prob = self.preconditioner_update_probability
-        if callable(update_prob):
-            update_prob = update_prob(self._prob_step)
-        do_update = self.rng.random() < update_prob
-        self._prob_step += 1
-
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
         min_ndim_triangular = group['min_ndim_triangular']
@@ -91,8 +78,8 @@ class ForeachPaLMPAdam(PSGDBase):
             if 'Q' not in state:
                 state['exp_avg'] = torch.zeros_like(g)
                 state['exp_avg_sq'] = torch.zeros_like(g)
-                Q, state["exprs"] = init_Q_exprs(p, precond_init_scale, max_size_triangular,
-                                                 min_ndim_triangular, memory_save_mode, dtype=q_dtype)
+                Q, state["exprs"] = init_Q_exprs(p, precond_init_scale, max_size_triangular, min_ndim_triangular,
+                                                 memory_save_mode, dtype=q_dtype)
                 state['Q'] = triu_to_line(Q) if store_triu_as_line else Q
 
             vals.append((p, g, state["Q"], state['exp_avg'], state['exp_avg_sq']))
@@ -106,11 +93,10 @@ class ForeachPaLMPAdam(PSGDBase):
         group["step"] += 1
 
         Q_triu = [line_to_triu(q) if store_triu_as_line else q for q in Q_list]
-        if do_update:
+        if self.should_update(group):
             for g, p, q_, q_orig in zip(grad_list, p_list, Q_triu, Q_list):
                 q32 = [promote(qq_) for qq_ in q_]
-                self.balance([g], [q32])
-                self.do_update([p], [g], [q32], precond_lr, [q_orig], store_triu_as_line=store_triu_as_line)
+                self.do_update(group, [p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
         torch._foreach_lerp_(exp_avg, grad_list, 1 - beta_debias(group['beta'], group['step']))
 
         beta2 = 1 - group['step'] ** -group['beta2_scale']

heavyball/psgd_kron.py

@@ -39,7 +39,7 @@ class ForeachPSGDKron(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: Optional[callable] = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32'):
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= beta < 1.0:
@@ -47,12 +47,8 @@ class ForeachPSGDKron(PSGDBase):
         if not 0.0 <= weight_decay:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
 
-        if preconditioner_update_probability is None:
-            preconditioner_update_probability = precond_update_prob_schedule()
         if clip_fn is None:
             clip_fn = lambda x: trust_region_clip_(x, 0.9, 1.5)
-        self.preconditioner_update_probability = preconditioner_update_probability
-        self.clip_fn = clip_fn
 
         defaults = dict(lr=lr, beta=beta, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
@@ -61,18 +57,10 @@ class ForeachPSGDKron(PSGDBase):
                         precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
                         step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
                         store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
-        super().__init__(params, defaults, foreach)
+        super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
-        self._prob_step = 0
 
     def _step(self, group):
-        # update preconditioners all together
-        update_prob = self.preconditioner_update_probability
-        if callable(update_prob):
-            update_prob = update_prob(self._prob_step)
-        do_update = self.rng.random() < update_prob
-        self._prob_step += 1
-
         momentum_into_precond_update = group.get("momentum_into_precond_update", True)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
@@ -114,10 +102,9 @@ class ForeachPSGDKron(PSGDBase):
             ea = exp_avg_list.pop(0)
             q = line_to_triu(q_orig) if store_triu_as_line else q_orig
 
-            if do_update:
+            if self.should_update(group):
                 q32 = [promote(q_) for q_ in q]
-                self.balance([ea if momentum_into_precond_update else g], [q32])
-                self.do_update([p], [g], [q32], precond_lr, [q_orig], store_triu_as_line=store_triu_as_line)
+                self.do_update(group,[p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
             set_(g, psgd_precond_grad(q, self.state_(p)["exprs"], ea))
 
         grad_list = self.clip_fn(grad_list)

heavyball/pure_psgd.py

@@ -5,7 +5,7 @@ Source available at https://github.com/evanatyourservice/kron_torch/blob/97a2b5e
 """
 
 import torch
-from heavyball.utils import copy_stochastic_list_
+from heavyball.utils import copy_stochastic_list_, identity
 
 from .utils import update_param_, warmup, psgd_precond_grad, init_Q_exprs, PSGDBase, precond_update_prob_schedule, \
     split_p_and_g_in_group, line_to_triu, triu_to_line, promote
@@ -38,18 +38,14 @@ class ForeachPurePSGD(PSGDBase):
                  max_size_triangular=2048, min_ndim_triangular=2, memory_save_mode=None,
                  momentum_into_precond_update=True, warmup_steps: int = 1, merge_dims: bool = False,
                  split: bool = False, clip_fn: callable = None, store_triu_as_line: bool = True,
-                 foreach: bool = True, q_dtype='float32'):
+                 foreach: bool = True, q_dtype='float32', stochastic_schedule: bool = True):
         if not 0.0 <= lr:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= weight_decay:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
 
-        if preconditioner_update_probability is None:
-            preconditioner_update_probability = precond_update_prob_schedule()
         if clip_fn is None:
-            clip_fn = lambda x: x
-        self.preconditioner_update_probability = preconditioner_update_probability
-        self.clip_fn = clip_fn
+            clip_fn = identity
 
         defaults = dict(lr=lr, weight_decay=weight_decay, max_size_triangular=max_size_triangular,
                         min_ndim_triangular=min_ndim_triangular, memory_save_mode=memory_save_mode,
@@ -58,18 +54,10 @@ class ForeachPurePSGD(PSGDBase):
                         precond_init_scale=1.0,  # precond init scale hardcoded to 1.0
                         step=0, warmup_steps=warmup_steps, merge_dims=merge_dims, split=split,
                         store_triu_as_line=store_triu_as_line, q_dtype=q_dtype)
-        super().__init__(params, defaults, foreach)
-
-        self._prob_step = 0
+        super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
         # update preconditioners all together
-        update_prob = self.preconditioner_update_probability
-        if callable(update_prob):
-            update_prob = update_prob(self._prob_step)
-        do_update = self.rng.random() < update_prob
-        self._prob_step += 1
-
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
         min_ndim_triangular = group['min_ndim_triangular']
@@ -105,10 +93,9 @@ class ForeachPurePSGD(PSGDBase):
             q_orig = Q_list.pop(0)
             q = line_to_triu(q_orig) if store_triu_as_line else q_orig
 
-            if do_update:
+            if self.should_update(group):
                 q32 = [promote(q_) for q_ in q]
-                self.balance([g], [q32])
-                self.do_update([p], [g], [q32], precond_lr, [q_orig], store_triu_as_line=store_triu_as_line)
+                self.do_update(group,[p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)
             psgd_precond_grad(q, self.state_(p)["exprs"], g, inplace=True)
 
         grad_list = self.clip_fn(grad_list)

heavyball/utils.py

@@ -668,7 +668,10 @@ def psgd_update_precond(Q, exprs, V, G, step, tiny):
 
         term2 += term1  # a + b
         term1 *= 2  # 2a
-        term1 -= term2  # 2a - (a + b) == a - b
+        if term1.dtype == term2.dtype:
+            term1 -= term2  # 2a - (a + b) == a - b
+        else:
+            term1 = term1 - term2
 
         term1 *= step
         norm = term2.norm(float('inf'))
@@ -790,21 +793,35 @@ def update_triu_(q_state, materialised):
 class PSGDBase(StatefulOptimizer):
     balance_probability: float = 0.01
 
-    def __init__(self, parameters, groups, foreach: bool = True):
-        super().__init__(parameters, groups, foreach)
+    def __init__(self, parameters, groups, foreach: bool, stochastic_schedule: bool, clip_fn,
+                 preconditioner_update_probability):
+        super().__init__(parameters, {**groups, 'stochastic_schedule': stochastic_schedule}, foreach)
         self.rng = random.Random(0x1923213)
         self._tiny = torch.finfo(torch.bfloat16).tiny
-
-    def balance(self, grad_list, Q_list):
-        if self.rng.random() > self.balance_probability:
-            return
-
-        for g, q in zip(grad_list, Q_list):
-            if g.dim() > 1:
-                psgd_balance_Q(q)
-
-    def do_update(self, p_list, grad_list, q_list, precond_lr, original_q: Optional[List] = None,
+        if clip_fn is None:
+            clip_fn = identity
+        if preconditioner_update_probability is None:
+            preconditioner_update_probability = precond_update_prob_schedule()
+        self.clip_fn = clip_fn
+        self.preconditioner_update_probability = preconditioner_update_probability
+
+    def should_update(self, group, prob: Optional[float] = None, name: str = 'cumulative_prob'):
+        group[f'{name}_prob_step'] = group.get(f'{name}_prob_step', 0) + 1
+        if prob is None:
+            prob = self.preconditioner_update_probability(group[f'{name}_prob_step'])
+        if group['stochastic_schedule']:
+            return self.rng.random() < prob
+        cumulative_prob = group.get(name, 0)
+        group[name] = cumulative_prob + prob
+        return int(group[name]) > int(cumulative_prob)
+
+    def do_update(self, group, p_list, grad_list, q_list, precond_lr, original_q: Optional[List] = None,
                   store_triu_as_line=False):
+        if self.should_update(group, self.balance_probability, 'balance_prob'):
+            for g, q in zip(grad_list, q_list):
+                if g.dim() > 1:
+                    psgd_balance_Q(q)
+
         for i, (p, grad, Q) in enumerate(zip(p_list, grad_list, q_list)):
             psgd_update_precond(Q, self.state_(p)["exprs"], torch.randn_like(grad), grad, precond_lr, self._tiny)
             if original_q:

{heavyball-0.17.2.dist-info → heavyball-0.18.0.dist-info}/METADATA

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

{heavyball-0.17.2.dist-info → heavyball-0.18.0.dist-info}/RECORD

@@ -1,24 +1,24 @@
 heavyball/__init__.py,sha256=iqP428JWwwx-XDOZ0nUdbCkOLEyfoqVyWZLQLAcwxaw,2214
-heavyball/cached_delayed_psgd_kron.py,sha256=DvjNNHzbnS-NDq965wve-VQ-ol7IFljYYGTuTwPHOhU,6971
-heavyball/cached_psgd_kron.py,sha256=xy3-yRKFUvRTstJb_asMVp-k-5Zuw_HyILPi7BsuMKQ,6974
-heavyball/delayed_psgd.py,sha256=rDDUj3miEn6HRJmKl-ZImsqkqBASSn8aC7MEV_06fzU,6017
+heavyball/cached_delayed_psgd_kron.py,sha256=JKNgqgT59Aa9evpCG--mOZcYC0qBqLntNU7uQYYWURM,6487
+heavyball/cached_psgd_kron.py,sha256=RCvXEiA-WwCERgndquaD-GU5L1wB4aeffkVDsbytaV4,6482
+heavyball/delayed_psgd.py,sha256=f3nX-6N6BQtZ3uIdZc7w2uJSphTJb4V1ocWpY35VdSU,5525
 heavyball/foreach_adamw.py,sha256=h_ar0ZZRM0q_wxuEkxEOEYe-2p-mB4OMgAHivrUnPl8,1777
 heavyball/foreach_adopt.py,sha256=ogOw2JjwEQNj7AKlweAphQFdMJ_GcMDm-RyDvEzugoc,1911
 heavyball/foreach_laprop.py,sha256=yGVmGqWiSw8Y2Xj70ndkR8ZMygakTB4_iRwV02Svkqg,1816
 heavyball/foreach_sfadamw.py,sha256=15-n6-lx4PAHYsKYmXbugxsR5MnqaPYy2vUudPRiitg,2087
 heavyball/foreach_soap.py,sha256=h6ptMch7oaynvu3eIJtWnVXypDA_5JDVm3Zb3PNEma0,4634
-heavyball/p_adam.py,sha256=F2b-xGNROi9VfX7isa3kffWePojpBl5BI1n464w4tGQ,6334
+heavyball/p_adam.py,sha256=lvPsG6tjV-NMY8lYTLVwmqTDz_BJg8SwbwQQrm55YlM,5849
 heavyball/palm_foreach_sfadamw.py,sha256=yvZbPyjDW8qd3r4qDXb6uTr5RozQ7JSDj4aYYRnKGLA,2248
 heavyball/palm_foreach_soap.py,sha256=g4hbiGRcti-J-a0SwAkP4ii5pU-aalsZH5bssyhroLk,5938
 heavyball/precond_schedule_foreach_soap.py,sha256=WLg5SzpJnKPZUvFyIvdwSZa1Umt5cpr3Kow_42orM-E,4863
 heavyball/precond_schedule_palm_foreach_soap.py,sha256=ammQrvRZFF-wc-wEiPEoFhS_7b8pdV61QfcLoQfimSo,6211
 heavyball/precond_schedule_sfpsoap.py,sha256=vq7jd302refKPa_9X2lkOTOtCCcTBVByPdojklrY8pA,6770
-heavyball/psgd_kron.py,sha256=2IpPj2TOExNGm8hSewi3er2GczJRNgC7r2J5yYSSA_0,5998
-heavyball/pure_psgd.py,sha256=uA7W9a3Qm1sxHQhtNxaUYrmE5x55lP5iJOKy_qT8XaQ,5341
+heavyball/psgd_kron.py,sha256=xBXO2rdzPtxcGjb5xDWPcFKixVoO9nwaaUUWckq1rBI,5466
+heavyball/pure_psgd.py,sha256=uMIQsslIbXuJ_uHJ9_bQD23mPv1IG0CZD6eA6W_AJ6g,4901
 heavyball/schedule_free_palm_foreach_soap.py,sha256=zkcikH5wWbzq4kOrmBjilvY3iWzuUddcv2HNEPKr3MI,6366
-heavyball/utils.py,sha256=8dhagAGj03D7kBEWOJmqsCjQKP069e1WwrzVp1JsBr8,29472
-heavyball-0.17.2.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.17.2.dist-info/METADATA,sha256=kCK9J8gg-6lj0qao7S-yDc7jsOzGxNJ6F4_JUFaiIR4,11810
-heavyball-0.17.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.17.2.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.17.2.dist-info/RECORD,,
+heavyball/utils.py,sha256=sojDNo94l-jAPDSnhEM5EI6D83SHG_hTnJOkdnDquSI,30492
+heavyball-0.18.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.18.0.dist-info/METADATA,sha256=_IjP5WGcuqV-ryQyy7pW9b5BRHxIJtsMFTRWuRDDs3o,11810
+heavyball-0.18.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.18.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.18.0.dist-info/RECORD,,