heavyball 0.24.0__py3-none-any.whl → 0.24.2__py3-none-any.whl

heavyball/utils.py

@@ -134,7 +134,7 @@ def dim_merger(grad, max_precond_dim, split: bool = False):
     if curr_shape > 1 or len(new_shape) == 0:
         new_shape.append(curr_shape)
 
-    new_grad = grad.view(new_shape)
+    new_grad = grad.reshape(new_shape)  # needs to be .reshape() due to channels_last
     if not split:
         return new_grad
 
@@ -1129,3 +1129,12 @@ 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 hook_optimizer_into_model(model, optimizer, *args, **kwargs):
+    def _step(p: Tensor, o: torch.optim.Optimizer):
+        o.step()
+        o.zero_grad()
+
+
+    for p in model.parameters():
+        p.register_post_accumulate_grad_hook(functools.partial(_step, o=optimizer([p], *args, **kwargs)))

{heavyball-0.24.0.dist-info → heavyball-0.24.2.dist-info}/METADATA

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

{heavyball-0.24.0.dist-info → heavyball-0.24.2.dist-info}/RECORD

@@ -16,9 +16,9 @@ heavyball/precond_schedule_sfpsoap.py,sha256=KUKdZzd336w24zPRcqwRatj7IVmd1Us0a_V
 heavyball/psgd_kron.py,sha256=PtTe6eR547Y-4CvgjpchgkQsr_kWr4AN-uY9L_JO_C8,6088
 heavyball/pure_psgd.py,sha256=344NdVNHwUFX3fU2R1S_Xh9SXAML3E4ryHr7xfMh9Cc,5076
 heavyball/schedule_free_palm_foreach_soap.py,sha256=KTQY37MZH7YnOSTLKY8uVySUXxWXbFVUA1QXN3iv8Ds,7244
-heavyball/utils.py,sha256=12DfrpBDiHAdFxN3cA3BA9tcailHw8wl5QTzEn4As98,39677
-heavyball-0.24.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.24.0.dist-info/METADATA,sha256=ZL_FTyrobNlPxhtgdVH6kZ9aD_jIq0SGALp7jo0BbiI,11926
-heavyball-0.24.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.24.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.24.0.dist-info/RECORD,,
+heavyball/utils.py,sha256=FglgQfiE206I07rql3qP-X2C1j0hY3N5VcQwKUh08aA,40025
+heavyball-0.24.2.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.24.2.dist-info/METADATA,sha256=lTThJQbW6qbnQqy9lGlTTOttJcX5vfQ_s6Cm0arqfC8,11926
+heavyball-0.24.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.24.2.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.24.2.dist-info/RECORD,,