heavyball 0.18.2__py3-none-any.whl → 0.18.3__py3-none-any.whl

heavyball/cached_delayed_psgd_kron.py

@@ -62,6 +62,7 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
+        should_update = self.should_update(group)
         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']
@@ -115,7 +116,7 @@ class ForeachCachedDelayedPSGDKron(PSGDBase):
 
             new = torch.einsum(self.state_(p)['cache_expr'], *cached_q, ea)
 
-            if self.should_update(group):
+            if should_update:
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig
                 q32 = [promote(q_) for q_ in q]
                 self.do_update(group, [p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig],

heavyball/cached_psgd_kron.py

@@ -8,8 +8,8 @@ from typing import Optional
 
 import torch
 
-from .utils import update_param_, warmup, init_Q_exprs, trust_region_clip_, PSGDBase, \
-    split_p_and_g_in_group, line_to_triu, triu_to_line, set_, einsum_base, promote
+from .utils import update_param_, warmup, init_Q_exprs, trust_region_clip_, PSGDBase, split_p_and_g_in_group, \
+    line_to_triu, triu_to_line, set_, einsum_base, promote
 
 
 class ForeachCachedPSGDKron(PSGDBase):
@@ -71,6 +71,7 @@ class ForeachCachedPSGDKron(PSGDBase):
         beta = group['beta']
         store_triu_as_line = group['store_triu_as_line']
         q_dtype = getattr(torch, group['q_dtype'])
+        should_update = self.should_update(group)
 
         vals = []
 
@@ -111,7 +112,7 @@ class ForeachCachedPSGDKron(PSGDBase):
             q_orig = Q_list.pop(0)
             ea = exp_avg_list.pop(0)
 
-            if self.should_update(group):
+            if should_update:
                 q = line_to_triu(q_orig) if store_triu_as_line else q_orig
                 q32 = [promote(q_) for q_ in q]
                 self.do_update(group, [p], [ea if momentum_into_precond_update else g], [q32], precond_lr, [q_orig],

heavyball/delayed_psgd.py

@@ -62,6 +62,7 @@ class ForeachDelayedPSGD(PSGDBase):
 
 
     def _step(self, group):
+        should_update = self.should_update(group)
         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']
@@ -103,7 +104,7 @@ 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 self.should_update(group):
+            if should_update:
                 q32 = [promote(q_) for q_ in q]
                 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)

heavyball/p_adam.py

@@ -61,6 +61,7 @@ class ForeachPaLMPAdam(PSGDBase):
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
+        should_update = self.should_update(group)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
         min_ndim_triangular = group['min_ndim_triangular']
@@ -94,7 +95,7 @@ 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 self.should_update(group):
+        if should_update:
             for g, p, q_, q_orig in zip(grad_list, p_list, Q_triu, Q_list):
                 q32 = [promote(qq_) for qq_ in q_]
                 self.do_update(group, [p], [g], [q32], precond_lr, [q_orig], store_triu_as_line)

heavyball/psgd_kron.py

@@ -60,6 +60,7 @@ class ForeachPSGDKron(PSGDBase):
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
+        should_update = self.should_update(group)
         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']
@@ -101,7 +102,7 @@ class ForeachPSGDKron(PSGDBase):
             ea = exp_avg_list.pop(0)
             q = line_to_triu(q_orig) if store_triu_as_line else q_orig
 
-            if self.should_update(group):
+            if should_update:
                 q32 = [promote(q_) for q_ in q]
                 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))

heavyball/pure_psgd.py

@@ -57,7 +57,7 @@ class ForeachPurePSGD(PSGDBase):
         super().__init__(params, defaults, foreach, stochastic_schedule, clip_fn, preconditioner_update_probability)
 
     def _step(self, group):
-        # update preconditioners all together
+        should_update = self.should_update(group)
         precond_init_scale = group['precond_init_scale']
         max_size_triangular = group['max_size_triangular']
         min_ndim_triangular = group['min_ndim_triangular']
@@ -93,7 +93,7 @@ class ForeachPurePSGD(PSGDBase):
             q_orig = Q_list.pop(0)
             q = line_to_triu(q_orig) if store_triu_as_line else q_orig
 
-            if self.should_update(group):
+            if group:
                 q32 = [promote(q_) for q_ in q]
                 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)

heavyball/utils.py

@@ -335,9 +335,7 @@ def promote(x):
 def min_dtype(xs: List[torch.Tensor]):
     dtypes = [x.dtype for x in xs]
     for d in (torch.float32, torch.bfloat16, torch.float16):
-        if all(d == x for x in dtypes):
-            return d
-        if all(d in (x, torch.float32, torch.float64) for x in dtypes):
+        if all(x in (d, torch.float32, torch.float64) for x in dtypes):
             return d
     return torch.float32
 

{heavyball-0.18.2.dist-info → heavyball-0.18.3.dist-info}/METADATA

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

{heavyball-0.18.2.dist-info → heavyball-0.18.3.dist-info}/RECORD

@@ -1,24 +1,24 @@
 heavyball/__init__.py,sha256=iqP428JWwwx-XDOZ0nUdbCkOLEyfoqVyWZLQLAcwxaw,2214
-heavyball/cached_delayed_psgd_kron.py,sha256=t8XXsl91lINY0iB5Cn5aQDyjLxN2itGiA97Pur4mEkY,6422
-heavyball/cached_psgd_kron.py,sha256=NnzgJB11xPfi5NrHl3OsQkgs-fxeR_tsHdfXeDiXxbE,6379
-heavyball/delayed_psgd.py,sha256=ylLNHglvjnkYAmJwcl1TtPA4PXKPaOv1YHVt0JVabMA,5551
+heavyball/cached_delayed_psgd_kron.py,sha256=PQAER6UgVh5l87DGRZrJ8CVP9UhyCG5wJD9rPLnj_G8,6460
+heavyball/cached_psgd_kron.py,sha256=GaeneBp0irksCSBIrJY4D_0hCpZ-uSRPMhqVX_a-og8,6417
+heavyball/delayed_psgd.py,sha256=fhBWFLTSl1S2gHWCeYak-STaXRwpC56sWZGLFMKFEJM,5589
 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=HnlOSR6fqOet0S4KavyU6zGtu7Tz7vUX18yJIgMnEBc,5845
+heavyball/p_adam.py,sha256=4zJDGJrpgUyVzr3GiELETFre4xr3-PE10OuAZj-jFM8,5883
 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=KhZnV5MpigAEfJfvYI7ApF1GQ8ZWWXl7g5nYueWKYDQ,5438
-heavyball/pure_psgd.py,sha256=qPQ46pp7DWyQ1afBin2bqFVhaRhjt7RjXm6VuM_2sxg,4851
+heavyball/psgd_kron.py,sha256=u46dorOUXx-do1IYeno2wj-6l1zYKMQQC-N2Zr2PzLI,5476
+heavyball/pure_psgd.py,sha256=iUy7mMKWxwNiVUMYrQ7SBnreu3t_XSbnhTW3a1yw4m0,4835
 heavyball/schedule_free_palm_foreach_soap.py,sha256=zkcikH5wWbzq4kOrmBjilvY3iWzuUddcv2HNEPKr3MI,6366
-heavyball/utils.py,sha256=LTkqF5hP6z9De4_jDLF0HQYUwz1MkJEOTdMOYyH5D0k,30426
-heavyball-0.18.2.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.18.2.dist-info/METADATA,sha256=H074F1jIGrFvWdTdeRmhGXoxJ5V0zH7lIEhZ-LSP3Mc,11810
-heavyball-0.18.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.18.2.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.18.2.dist-info/RECORD,,
+heavyball/utils.py,sha256=qs_WfzJdS-3XyEuw-m6mWMEeR95r7bGFVC8wWCHtD48,30365
+heavyball-0.18.3.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.18.3.dist-info/METADATA,sha256=Cx8LM2g3BFOk8WJH3B8ve8kQ7HghMCIRLggdJp37x4g,11810
+heavyball-0.18.3.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.18.3.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.18.3.dist-info/RECORD,,