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

heavyball/utils.py

@@ -479,7 +479,7 @@ def copy_stochastic_list_(target: List[torch.Tensor], source: List[torch.Tensor]
         copy_stochastic_(t, s)
 
 
-@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=True)
 def _compilable_copy_stochastic_(target: torch.Tensor, source: torch.Tensor):
     """Taken as-is from https://github.com/pytorch/pytorch/issues/120376#issuecomment-1974828905"""
     # create a random 16 bit integer
@@ -815,20 +815,21 @@ class PSGDBase(StatefulOptimizer):
 
     def do_update(self, group, p_list, grad_list, q_list, precond_lr, original_q: Optional[List] = None,
                   store_triu_as_line=False):
-        for i, (p, grad, Q) in enumerate(zip(p_list, grad_list, q_list)):
+        if original_q:
+            if store_triu_as_line:
+                update_fn = update_triu_
+            else:
+                update_fn = copy_stochastic_list_
+        else:
+            update_fn = lambda x, y: None
+        for i, (p, grad, Q, oq) in enumerate(zip(p_list, grad_list, q_list, original_q)):
             psgd_update_precond(Q, self.state_(p)["exprs"], torch.randn_like(grad), grad, precond_lr, self._tiny)
+            update_fn(oq, Q)
 
-        for g, q in zip(grad_list, q_list):
+        for g, q in zip(grad_list, original_q if original_q else q_list):
             if g.dim() > 1:
                 psgd_balance_Q(q)
 
-        if original_q:
-            for q in q_list:
-                if store_triu_as_line:
-                    update_triu_(original_q[i], Q)
-                else:
-                    copy_stochastic_list_(original_q[i], Q)
-
 
 def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.001, flat_start=250):
     """Anneal preconditioner update probability during beginning of training.

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

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

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

@@ -16,9 +16,9 @@ heavyball/precond_schedule_sfpsoap.py,sha256=vq7jd302refKPa_9X2lkOTOtCCcTBVByPdo
 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=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,,
+heavyball/utils.py,sha256=U8BX11BAhUCH_JO_tnVy4JLSkfDSBHFi_a5s8Pvsf-s,30437
+heavyball-0.18.4.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.18.4.dist-info/METADATA,sha256=kBGWYUMlF21_Cg2qtcHa4ajfzB3xlm5UUb24S3mUgKI,11810
+heavyball-0.18.4.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.18.4.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.18.4.dist-info/RECORD,,