heavyball 0.21.8__py3-none-any.whl → 0.23.0__py3-none-any.whl

heavyball/utils.py

@@ -7,6 +7,7 @@ from typing import List, Optional, Tuple, Callable, Union
 
 import numpy as np
 import torch
+from torch import Tensor
 from torch.backends import cudnn, opt_einsum
 from torch.utils._pytree import tree_map
 
@@ -39,15 +40,14 @@ def warmup(lr: float, step: int, warmup_steps: int):
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
-def _compilable_schedule_free_(p, z, ckp1, grad, lr, beta1):
-    p32 = promote(p)
-    z32 = promote(z)
-    p32.lerp_(end=z32, weight=ckp1)
-    p32.add_(grad, alpha=lr * (beta1 * (1 - ckp1) - 1))
-    copy_stochastic_(p, p32)
-
-    z32.add_(grad, alpha=-lr)
-    copy_stochastic_(z, z32)
+def _compilable_schedule_free_(p: List[Tensor], z: List[Tensor], ckp1: Tensor, grad: List[Tensor], lr: Tensor, beta1: Tensor):
+    p32, z32, g32 = [promote(x) for x in (p, z, grad)]
+    for p_, z_, g_ in zip(p32, z32, g32):
+        p_.lerp_(z_, ckp1)
+        p_.add_(g_, alpha=lr * (beta1 * (1 - ckp1) - 1))
+        z_.add(g_, alpha=-lr)
+    copy_stochastic_list_(p, p32)
+    copy_stochastic_list_(z, z32)
 
 
 def get_ckp1(lr, weight_lr_power, weight_sum, r, step):
@@ -61,8 +61,8 @@ def get_ckp1(lr, weight_lr_power, weight_sum, r, step):
     return ckp1, weight_sum
 
 
-def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1: float, parameters: List[torch.Tensor],
-                   z: List[torch.Tensor], grad: list[torch.Tensor], r: float = 0.0, step: int = 0):
+def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1: float, parameters: List[Tensor],
+                   z: List[Tensor], grad: list[Tensor], r: float = 0.0, step: int = 0):
     weight = lr ** weight_lr_power * max(step, 1) ** r
     weight_sum = weight_sum + weight
 
@@ -73,10 +73,8 @@ def schedule_free_(lr: float, weight_lr_power: float, weight_sum: float, beta1:
 
     # These operations update y in-place,
     # without computing x explicitly.
-    lr_tensor = torch.empty((), dtype=torch.float32, device=parameters[0].device).fill_(lr)
-    ckp1_tensor = torch.empty((), dtype=torch.float32, device=parameters[0].device).fill_(ckp1)
-    for p, z_, g in zip(parameters, z, grad):
-        _compilable_schedule_free_(p, z_, ckp1_tensor, g, lr_tensor, beta1)
+    lr, ckp1 = scalar_guard(lr, parameters[0]), scalar_guard(ckp1, parameters[0])
+    _compilable_schedule_free_(parameters, z, ckp1, grad, lr, beta1)
     return weight_sum
 
 
@@ -142,19 +140,25 @@ def beta_debias(beta, step):
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
-def exp_avg_sq_(state, grad, beta2, eps, out=None):
-    if isinstance(state, torch.Tensor):
-        state.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
-        return torch.sqrt(state, out=out).clamp_(min=eps)
-
+def _compilable_exp_avg_sq_(state: List[Tensor], grad: List[Tensor], beta2: Tensor, eps: Tensor, out: List[Optional[Tensor]]):
     torch._foreach_mul_(state, beta2)
     [s.addcmul_(g, g, value=1 - beta2) for s, g in zip(state, grad)]
     denom = torch._foreach_sqrt(state)
-    torch._foreach_maximum_(denom, eps)
-    return denom
+    [denom.clamp_(min=eps) for denom in denom]
+    if out[0] is None:
+        return denom
 
+    copy_stochastic_list_(out, denom)
+    return out
 
-def adaptive_gradient_clipping_(parameters: List[torch.Tensor], gradients: List[torch.Tensor], clip_val: float,
+
+def exp_avg_sq_(state, grad, beta2, eps, out=None):
+    state, grad, out = list_guard(state), list_guard(grad), list_guard(out)
+    beta2, eps = scalar_guard(beta2, state[0]), scalar_guard(eps, state[0])
+    return _compilable_exp_avg_sq_(state, grad, beta2, eps, out)
+
+
+def adaptive_gradient_clipping_(parameters: List[Tensor], gradients: List[Tensor], clip_val: float,
                                 minimum: float = 1e-3, eps: float = 1e-8):
     if clip_val <= 0:
         return
@@ -168,12 +172,19 @@ def adaptive_gradient_clipping_(parameters: List[torch.Tensor], gradients: List[
     torch._foreach_mul_(gradients, p_norm)
 
 
-def set_(dst: torch.Tensor, src: torch.Tensor):
-    if not torch.compiler.is_compiling() and src.data_ptr() == dst.data_ptr():
+def is_compiling():
+    try:
+        return torch.compiler.is_compiling()
+    except AttributeError:
+        return True
+
+
+def set_(dst: Tensor, src: Tensor):
+    if not is_compiling() and src.data_ptr() == dst.data_ptr():
         return
     if src.shape != dst.shape:
         src = src.reshape_as(dst)
-    if not torch.compiler.is_compiling() and src.is_contiguous() and dst.is_contiguous() and src.dtype == dst.dtype:
+    if not is_compiling() and src.is_contiguous() and dst.is_contiguous() and src.dtype == dst.dtype:
         dst.set_(src)
     else:
         dst.copy_(src)
@@ -329,7 +340,7 @@ def get_orthogonal_matrix(mat):
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
-def _compilable_stochastic_lerp_(x: List[torch.Tensor], y: List[torch.Tensor], a: Union[float, int, torch.Tensor]):
+def _compilable_stochastic_lerp_(x: List[Tensor], y: List[Tensor], a: Union[float, int, Tensor]):
     for x_, y_ in zip(x, y):
         x32 = promote(x_)
         y32 = promote(y_)
@@ -337,14 +348,28 @@ def _compilable_stochastic_lerp_(x: List[torch.Tensor], y: List[torch.Tensor], a
         copy_stochastic_(x_, x32)
 
 
-def stochastic_lerp_(x: List[torch.Tensor], y: List[torch.Tensor], a: Union[float, int, torch.Tensor]):
-    if not isinstance(a, torch.Tensor):
-        a = torch.empty((), dtype=torch.float32, device=x[0].device).fill_(a)
+def stochastic_lerp_(x: List[Tensor], y: List[Tensor], a: Union[float, int, Tensor]):
+    x, y = list_guard(x), list_guard(y)
+    a = scalar_guard(a, x[0])
     _compilable_stochastic_lerp_(x, y, a)
 
 
+def list_guard(x):
+    if isinstance(x, (list, tuple)):
+        return x
+    return [x]
+
+
+def scalar_guard(x, ref):
+    if isinstance(x, float):
+        return torch.empty((), dtype=torch.float32, device=ref.device).fill_(x)
+    if isinstance(x, int):
+        return torch.empty((), dtype=torch.int64, device=ref.device).fill_(x)
+    return x
+
+
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
-def _compilable_stochastic_add_(x: List[torch.Tensor], y: List[torch.Tensor], alpha: Union[float, int, torch.Tensor]):
+def _compilable_stochastic_add_(x: List[Tensor], y: List[Tensor], alpha: Union[float, int, Tensor]):
     for x_, y_ in zip(x, y):
         x32 = promote(x_)
         y32 = promote(y_)
@@ -352,9 +377,9 @@ def _compilable_stochastic_add_(x: List[torch.Tensor], y: List[torch.Tensor], al
         copy_stochastic_(x_, x32)
 
 
-def stochastic_add_(x: List[torch.Tensor], y: List[torch.Tensor], alpha: Union[float, int, torch.Tensor]):
-    if not isinstance(alpha, torch.Tensor):
-        alpha = torch.empty((), dtype=torch.float32, device=x[0].device).fill_(alpha)
+def stochastic_add_(x: List[Tensor], y: List[Tensor], alpha: Union[float, int, Tensor]):
+    x, y = list_guard(x), list_guard(y)
+    alpha = scalar_guard(alpha, x[0])
     _compilable_stochastic_add_(x, y, alpha)
 
 
@@ -376,12 +401,12 @@ def compute_ggt(grad, GG, max_precond_dim, precondition_1d, beta):
 def promote(x):
     if isinstance(x, torch.dtype) and x in (torch.bfloat16, torch.float16):
         return torch.float32
-    if isinstance(x, torch.Tensor) and x.dtype in (torch.bfloat16, torch.float16):
+    if isinstance(x, Tensor) and x.dtype in (torch.bfloat16, torch.float16):
         return x.float()
     return x
 
 
-def min_dtype(xs: List[torch.Tensor]):
+def min_dtype(xs: List[Tensor]):
     dtypes = [x.dtype for x in xs]
     for d in (torch.float32, torch.bfloat16, torch.float16):
         if all(x in (d, torch.float32, torch.float64) for x in dtypes):
@@ -447,7 +472,7 @@ class StatefulOptimizer(torch.optim.Optimizer):
         self.fake_groups = {}
         self.use_ema = use_ema
 
-    def key(self, param: torch.Tensor):
+    def key(self, param: Tensor):
         return (param.data_ptr(), tuple(param.shape))
 
     def get_groups(self, group):
@@ -460,19 +485,56 @@ class StatefulOptimizer(torch.optim.Optimizer):
 
         return [self.fake_groups[self.key(p)] for p in group['params']]
 
-    def state_(self, arg: torch.Tensor):
+    def state_(self, arg: Tensor):
         return self.state[self.key(arg)]
 
+    def mars_correct_list(self, group, p_list, g_list, mars_gamma, beta):
+        for p, g in zip(p_list, g_list):
+            state = self.state_(p)
+            if 'mars_old_grad' not in state:
+                state['mars_old_grad'] = torch.zeros_like(g)
+        old_gs = [self.state_(p)['mars_old_grad'] for p in p_list]
+        mars_correction(g_list, old_gs, mars_gamma, beta)
+
+    def split_p_and_g_in_group(self, group: dict, skip_none: bool = True, should_promote: bool = True,
+                               beta1: float = -1.0):
+        for p in group["params"]:
+            if skip_none and p.grad is None:
+                continue
+
+            if p.grad is None:
+                grad = None
+            else:
+                if should_promote:
+                    grad = promote(p.grad)
+                else:
+                    grad = p.grad
+                if beta1 >= 0 and group.get('mars', False):
+                    self.mars_correct_list(group, [p], [grad], group['mars_gamma'], beta1)
+
+                p.grad = None
+
+            p_views = merge_group(group, p)
+            if grad is not None:
+                grad = merge_group(group, grad)
+            if isinstance(p_views, Tensor):
+                yield p_views, grad
+                continue
+            if grad is None:
+                yield from zip(p_views, [None] * len(p_views))
+                continue
+            yield from zip(p_views, grad)
+
     def state_size(self) -> int:
         total_bytes = 0
 
         def _add(x):
             nonlocal total_bytes
-            if isinstance(x, torch.Tensor):
+            if isinstance(x, Tensor):
                 total_bytes += x.numel() * x.element_size()
 
         for group in self.param_groups:
-            for p, _ in split_p_and_g_in_group(group, skip_none=False):
+            for p, _ in self.split_p_and_g_in_group(group, skip_none=False):
                 tree_map(_add, self.state_(p))
         return total_bytes
 
@@ -576,13 +638,14 @@ class ScheduleFree(StatefulOptimizer):
         raise NotImplementedError
 
 
-def copy_stochastic_list_(target: List[torch.Tensor], source: List[torch.Tensor]):
+def copy_stochastic_list_(target: List[Tensor], source: List[Tensor]):
     for t, s in zip(target, source):
         copy_stochastic_(t, s)
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
-def _compilable_exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step):
+def _compilable_exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor],
+                         grad_projected: List[Tensor], beta1: Tensor, beta2: Tensor, step: Tensor):
     beta1 = beta_debias(beta1, step)
     beta2 = beta_debias(beta2, step)
 
@@ -595,21 +658,17 @@ def _compilable_exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2
     return denom
 
 
-def exp_avg_(exp_avg: List[torch.Tensor], exp_avg_sq: List[torch.Tensor], grad: List[torch.Tensor],
-             grad_projected: List[torch.Tensor], beta1: float, beta2: float, step: int):
-    if isinstance(beta1, float):
-        beta1 = torch.empty((), dtype=torch.float32, device=exp_avg[0].device).fill_(beta1)
-    if isinstance(beta2, float):
-        beta2 = torch.empty((), dtype=torch.float32, device=exp_avg[0].device).fill_(beta2)
-    if isinstance(step, int):
-        step = torch.empty((), dtype=torch.int32, device=exp_avg[0].device).fill_(step)
+def exp_avg_(exp_avg: List[Tensor], exp_avg_sq: List[Tensor], grad: List[Tensor], grad_projected: List[Tensor],
+             beta1: float, beta2: float, step: int):
+    exp_avg, exp_avg_sq, grad, grad_projected = list_guard(exp_avg), list_guard(exp_avg_sq), list_guard(
+        grad), list_guard(grad_projected)
+    beta1, beta, step = scalar_guard(beta1, exp_avg[0]), scalar_guard(beta2, exp_avg[0]), scalar_guard(step, exp_avg[0])
     denom = _compilable_exp_avg_(exp_avg, exp_avg_sq, grad, grad_projected, beta1, beta2, step)
     return denom
 
 
-# this can be dynamic for most optimizers - just not for PSGD. So, it's disabled for all
-@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True)
-def _compilable_copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+def _compilable_copy_stochastic_(target: Tensor, source: Tensor):
     """Taken as-is from https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905"""
     # create a random 16 bit integer
     result = torch.randint_like(source, dtype=torch.int32, low=0, high=(1 << 16))
@@ -624,8 +683,8 @@ def _compilable_copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
     target.copy_(result.view(dtype=torch.float32))
 
 
-def copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
-    if not torch.compiler.is_compiling() and target.data_ptr() == source.data_ptr():
+def copy_stochastic_(target: Tensor, source: Tensor):
+    if not is_compiling() and target.data_ptr() == source.data_ptr():
         return
     if target.dtype != torch.bfloat16 or source.dtype not in (torch.float16, torch.float32, torch.float64):
         set_(target, source)
@@ -633,26 +692,31 @@ def copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
-def _compilable_update_(p, u, decay, add_fn, lr):
+def _compilable_update_(p: List[Tensor], u: List[Tensor], decay: Tensor, add_fn: callable, 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]]
 
     if decay > 0:
         torch._foreach_mul_(p32, 1 - decay * lr)
 
-    for p32_, u32_ in zip(p32, u32):  # lr is data-dependent -> can't compile a foreach
-        if add_fn is None:
-            p32_.add_(u32_, alpha=lr)
-        else:
-            add_fn(p32_, u32_, lr)
+    for p32_, u32_, g_ in zip(p32, u32, g):  # lr is data-dependent -> can't compile a foreach
+        if caution:
+            _compilable_cautioning_(promote(g_), u32_)
+        add_fn(p32_, u32_, lr)
 
     copy_stochastic_list_(p, p32)
 
 
-def update_param_(param: List[torch.Tensor], update: List[torch.Tensor], lr: float, decay: float,
-                  add_fn: callable = None):
-    lr_tensor = torch.empty((), dtype=torch.float32, device=param[0].device).fill_(lr)
-    _compilable_update_(param, update, decay, add_fn, lr_tensor)
+def update_param_(param: List[Tensor], update: List[Tensor], lr: float, decay: float, add_fn: callable = None,
+                  caution: bool = False, grad: List[Tensor] = None):
+    param, update, grad = list_guard(param), list_guard(update), list_guard(grad)
+    lr = scalar_guard(lr, param[0])
+    if not caution:
+        grad = [None] * len(param)
+    if add_fn is None:
+        add_fn = stochastic_add_
+    _compilable_update_(param, update, decay, add_fn, lr, caution, grad)
 
 
 def precond_schedule(step, precond_scheduler, rng):
@@ -822,14 +886,14 @@ def psgd_update_precond(Q, exprs, G, precond_lr, tiny, oq, store_triu_as_line):
 
 
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
-def psgd_precond_grad(Q, exprs, G, inplace: bool = False):
+def psgd_precond_grad(inplace: bool, exprs: str, grad: Tensor, *preconds: Tensor):
     """Precondition gradient G with preconditioner Q."""
-    md = min_dtype(Q)
-    out = torch.einsum(exprs[-1], *[q.conj().to(md) for q in Q], *[q.to(md) for q in Q], G.to(md))
+    md = min_dtype(preconds)
+    out = torch.einsum(exprs, *[q.conj().to(md) for q in preconds], *[q.to(md) for q in preconds], grad.to(md))
     if inplace:
-        set_(G, out)
-        return G
-    return out.to(G.dtype)
+        set_(grad, out)
+        return grad
+    return out.to(grad.dtype)
 
 
 def norm_clip_(x, scale=None):
@@ -892,7 +956,7 @@ def trust_region_clip_(grad, lerp: float = 0.9, scale: float = 1.5):
 
 
 @decorator
-def triu_to_line(Q_list: List[torch.Tensor]):
+def triu_to_line(Q_list: List[Tensor]):
     out = []
     for q in Q_list:
         if q.dim() < 2:
@@ -909,7 +973,7 @@ def _triu_shape(numel):
 
 
 @decorator
-def line_to_triu(Q_list: List[Tuple[Optional[List[int]], torch.Tensor]]):
+def line_to_triu(Q_list: List[Tuple[Optional[List[int]], Tensor]]):
     new = []
     for shape, q in Q_list:
         if shape is not None:
@@ -965,18 +1029,45 @@ class PSGDBase(StatefulOptimizer):
                         psgd_balance_Q(q)
 
 
+# TODO: Figure out why this sometimes crashes
 @torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
-def _compilable_precond_grad_cached_(cached_q, ea, expr, param, lr, weight_decay, clip_fn):
+def _compilable_precond_grad_cached_(ea: Tensor, expr: str, param: Tensor, lr: Tensor, weight_decay: Tensor,
+                                     clip_fn: callable, caution: bool, grad: Optional[Tensor], *cached_q: Tensor):
     md = min_dtype(cached_q + [ea])
     new = torch.einsum(expr, *[c_.to(md) for c_ in cached_q], ea.to(md)).to(torch.float32)
-    update_param_([param], clip_fn([new]), lr, weight_decay)
+    update_param_([param], clip_fn([new]), lr, weight_decay, caution=caution, grad=grad)
+
+
+def precond_grad_cached_(cached_q: List[Tensor], ea: Tensor, expr: str, param: Tensor, lr: float, weight_decay: float,
+                         clip_fn, caution, grad):
+    lr = scalar_guard(lr, param)
+    _compilable_precond_grad_cached_(ea, expr, param, lr, weight_decay, clip_fn, caution, grad, *cached_q)
+
+
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+def _compilable_mars_correction_(g: Tensor, old_g: Tensor, a: Tensor):
+    g_copy = [g_.clone() for g_ in g]
+    _compilable_stochastic_lerp_(g, old_g, a)
+    copy_stochastic_list_(old_g, g_copy)
+
+
+def mars_correction(g, old_g, beta1, gamma):
+    a = -gamma * beta1 / (1 - beta1)
+    g, old_g = list_guard(g), list_guard(old_g)
+    a = scalar_guard(a, g[0])
+    _compilable_mars_correction_(g, old_g, a)
 
 
-def precond_grad_cached_(cached_q: List[torch.Tensor], ea: torch.Tensor, expr: str, param: torch.Tensor, lr: float,
-                         weight_decay: float, clip_fn):
-    if isinstance(lr, float):
-        lr = torch.empty((), dtype=torch.float32, device=param.device).fill_(lr)
-    _compilable_precond_grad_cached_(cached_q, ea, expr, param, lr, weight_decay, clip_fn)
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+def _compilable_cautioning_(g: Tensor, update: Tensor):
+    mask = (g * update) > 0
+    update.masked_fill_(~mask, 0)
+    scale = mask.numel() / mask.sum().clamp(min=1)
+    update.mul_(scale)
+
+
+def caution(g, update):
+    _compilable_cautioning_(g, update)
 
 
 def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.001, flat_start=250):
@@ -1013,29 +1104,3 @@ def merge_group(group, *tensors):
         append_or_extend(out, dim_merger(t, group['max_size_triangular'] if 'max_size_triangular' in group else group[
             'max_precond_dim'], group.get('split', False)))
     return out
-
-
-def split_p_and_g_in_group(group: dict, skip_none: bool = True, should_promote: bool = True):
-    for p in group["params"]:
-        if skip_none and p.grad is None:
-            continue
-
-        if p.grad is None:
-            grad = None
-        else:
-            if should_promote:
-                grad = promote(p.grad)
-            else:
-                grad = p.grad
-            p.grad = None
-
-        p_views = merge_group(group, p)
-        if grad is not None:
-            grad = merge_group(group, grad)
-        if isinstance(p_views, torch.Tensor):
-            yield p_views, grad
-            continue
-        if grad is None:
-            yield from zip(p_views, [None] * len(p_views))
-            continue
-        yield from zip(p_views, grad)

{heavyball-0.21.8.dist-info → heavyball-0.23.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.21.8
+Version: 0.23.0
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler
@@ -32,7 +32,7 @@ A simple package of efficient optimizers
 The goal is not to thrive for completeness, full maintenance or abstraction, but instead to provide a simple
 largely static alternative to `torch.optim` with more and better optimizers.
 
-Currently (2024-11-22, 0.21.0), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
+Currently (2024-11-26, 0.22.1), the recommended stable optimizer is `PrecondSchedulePaLMSOAP` (see below). The
 recommended experimental optimizer is `DelayedPSGDKron` ([tuning guide](docs/psgd_efficiency.md)).
 
 ## Features

heavyball-0.23.0.dist-info/RECORD

@@ -0,0 +1,24 @@
+heavyball/__init__.py,sha256=icHYN-MGsmHkLUlHCMcZkOlwY7GT63_ayR_a5iPKmzM,2226
+heavyball/cached_delayed_psgd_kron.py,sha256=n3wIOhrop0Ls4MZ0kXpwGuImp1jzPs6VGdxIlPyoYdQ,6827
+heavyball/cached_psgd_kron.py,sha256=KCLsfvj9qh_2FNwRTdWM3zjnt2oGHfsf4Y341rPcceI,6778
+heavyball/delayed_psgd.py,sha256=z_Y1eYr2upVt_FsyCIv91yTFJY6yqvHsI8S2mOpqdv8,6334
+heavyball/foreach_adamw.py,sha256=uawSbGGUD2E1RtcwspP83yQNElERdGX-diqCI5e8FqE,2825
+heavyball/foreach_adopt.py,sha256=DFEaPswVzdHcbxC-mirsf_okM_HR6r34PDUTty5CrUE,3547
+heavyball/foreach_laprop.py,sha256=J4Vms0nAOMh3GQtAOPyrYOe5WtpzokVv25b9oDnwc2A,2833
+heavyball/foreach_sfadamw.py,sha256=HWbLekY5BloHDIgrN2J0a7IolZCt8Ah2xkLAU_-5oSc,3079
+heavyball/foreach_soap.py,sha256=7B_dP2Hm_xqwpBQiPYkv_c6eoRnU1dV2VZfvSoa4uJ8,4729
+heavyball/p_adam.py,sha256=8BlZ6YoaDXawMiRbCxo0Kd5_0-pAn0MQIhL0LHNaRBs,6315
+heavyball/palm_foreach_sfadamw.py,sha256=E8raxrBIkSmTEGFzwnfWxKwDJjBQE2vdsmyqfc8aL_A,3375
+heavyball/palm_foreach_soap.py,sha256=IknGm_CzrqDIFEoCkejxjoZ4sfIy6RSoInqlMUOYLB4,6156
+heavyball/precond_schedule_foreach_soap.py,sha256=bJ2ifPFa8zEP9GO8eBpqZzsmP7p_iQkkCkllNeEMHPU,4892
+heavyball/precond_schedule_palm_foreach_soap.py,sha256=4dT9f134-Faq2KuCMCHzMtrkMO-es5p_DYS1of5yF-s,6428
+heavyball/precond_schedule_sfpsoap.py,sha256=FOR-axwlkSN7IHZWYYUVFfjSFCLxc_NdiTlb-n5gmgs,7530
+heavyball/psgd_kron.py,sha256=4eiGPXAFjvGIXLdiai1UJfAvTozAV1TXaE9UGkE4BLc,6051
+heavyball/pure_psgd.py,sha256=344NdVNHwUFX3fU2R1S_Xh9SXAML3E4ryHr7xfMh9Cc,5076
+heavyball/schedule_free_palm_foreach_soap.py,sha256=0WT_gvTKymqLQzYT6ewDgCmpDq-HgMAewipw1QvyQYA,7267
+heavyball/utils.py,sha256=AZlY8dfM0d-C0FXBCJHTJOOoi3RjkMJ-XhU25aBN878,39521
+heavyball-0.23.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.23.0.dist-info/METADATA,sha256=3IBUhXA7VJT9GQh460OznCAcIqCG_Mv5Q7HZO8FQ40w,11926
+heavyball-0.23.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.23.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.23.0.dist-info/RECORD,,

heavyball-0.21.8.dist-info/RECORD

@@ -1,24 +0,0 @@
-heavyball/__init__.py,sha256=iqP428JWwwx-XDOZ0nUdbCkOLEyfoqVyWZLQLAcwxaw,2214
-heavyball/cached_delayed_psgd_kron.py,sha256=Nyxl-G-o6greKwDN-vLiw5W02GXO2LRvknc0OzvzFnE,6674
-heavyball/cached_psgd_kron.py,sha256=HzD6se0AYb-W5hpydUxcR9uqrpe_54PBwgL1VWX3DHU,6592
-heavyball/delayed_psgd.py,sha256=m4c-OvcLMrRxSAPYs2l6Up21uCyF2kvHvpcnfe3nzGs,6212
-heavyball/foreach_adamw.py,sha256=Rb5U80cgUcEqlEbUU250UTWdoqA7nyiqkV5w1U4bWX4,2445
-heavyball/foreach_adopt.py,sha256=ecdi1fKg9i087OGjtKWVbE_DD6Yf4pvpzv4ELCcusvQ,3211
-heavyball/foreach_laprop.py,sha256=vi6C_gfjXxw5uN0KHgzxI9itUI1dcgOf3ufoO_VVMp0,2471
-heavyball/foreach_sfadamw.py,sha256=rLZORmCIMu9G09FdDgMSiI6pNq34IVoxsPVWtmeDdbQ,2753
-heavyball/foreach_soap.py,sha256=4mWSMWYTdjgiXiboI5DwdigecruDtNGKylGAFAVhCRA,4562
-heavyball/p_adam.py,sha256=Xyxsavwtw-t0OyTHitYQXZSmF9UJlMDzDAURge-MbbQ,6047
-heavyball/palm_foreach_sfadamw.py,sha256=JbNrcoquBGGUI5XNMFouDjpNurVHUW9DbX1A3tSrtno,3025
-heavyball/palm_foreach_soap.py,sha256=GzAwM8kOt1X0QCmUZDTdHwPxbJwjH8ic43dyAK5BYCA,6015
-heavyball/precond_schedule_foreach_soap.py,sha256=HcObXLfSNN_lKNb4nmC6tkdHcqDIMNX6hILpHKScqLc,4744
-heavyball/precond_schedule_palm_foreach_soap.py,sha256=xZ7CJvIfdu2RNAZt2g1S7Xb0Jyy1hNC4MowOFU3nWkk,6283
-heavyball/precond_schedule_sfpsoap.py,sha256=PNneiOkrRyV1yIZn91lPmYofd1_OiLqJTDy75RLpXJk,7273
-heavyball/psgd_kron.py,sha256=RSLJi5FnSmYjvYBufNDClnYRm-eN_Kpa1Ar2tNP6-X0,5824
-heavyball/pure_psgd.py,sha256=LZK0qmvZkBF8g00evaVLtW-sIUJmdoxag1K7O26AqEo,4820
-heavyball/schedule_free_palm_foreach_soap.py,sha256=_AqrnChY6iDlQUkF2YUxS7eLjSWCIuvEUOHvMHVM1yY,6873
-heavyball/utils.py,sha256=xTDZEt2_DM57EYnJkRq7d7scTnro4eKPdMtEwPdLy-c,37218
-heavyball-0.21.8.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.21.8.dist-info/METADATA,sha256=nLyxHlENmhAGyU9GManYKKJJTykhsAMt7hkJNXPu_YY,11926
-heavyball-0.21.8.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.21.8.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.21.8.dist-info/RECORD,,