heavyball 0.20.0__py3-none-any.whl → 0.20.1__py3-none-any.whl

heavyball/utils.py

@@ -721,7 +721,7 @@ def init_Q_exprs(t, scale, max_size, min_ndim_triangular, memory_save_mode, dtyp
     return [Q, (exprA, tuple(exprGs), exprP)]
 
 
-@torch.compile(mode='max-autotune-no-cudagraphs', fullgraph=True, dynamic=False)
+@decorator
 def psgd_balance_Q(Q_in):
     norms = torch.stack([q.norm(float("inf")) for q in Q_in])
     geometric_mean = norms.log().mean().exp()
@@ -734,8 +734,7 @@ def psgd_calc_A_and_conjB(exprA, G, Q):
     A = torch.einsum(exprA, *[q.to(md) for q in Q], G.to(md)).to(G.dtype)
     order = G.dim()
     p = list(range(order))
-    V = torch.randn_like(G, dtype=promote(G.dtype))
-    conjB = torch.permute(V, p[1:] + p[:1])
+    conjB = torch.randn_like(G.shape[1:] + G.shape[:1], dtype=promote(G.dtype), device=G.device)
     Q = [promote(q) for q in Q]
     for i, q in enumerate(Q):
         if q.dim() <= 1:
@@ -755,21 +754,13 @@ def psgd_calc_A_and_conjB(exprA, G, Q):
 def psgd_lb(A, max_abs):
     A /= max_abs
     a0 = torch.einsum('ij,ij->j', A, A)
-    a1 = torch.einsum('ij,ij->i', A, A)
-    value0 = torch.max(a0)
-    value1 = torch.max(a1)
     i = torch.argmax(a0)
-    j = torch.argmax(a1)
 
-    comp = value0 > value1
-    x = torch.cond(comp, lambda a: torch.index_select(a, 1, i).flatten().contiguous(),  #
-                   lambda a: torch.index_select(a, 0, j).flatten().contiguous(), (A,))
+    x = torch.index_select(a, 1, i).flatten().contiguous()
 
-    x = torch.cond(comp, lambda x_, a: torch.einsum('i,ij->j', x_, a), lambda x_, a: torch.einsum('i,ji->j', x_, a),
-                   (x, A,))
+    x = torch.einsum('i,ij->j', x_, a)
     x /= x.norm()
-    x = torch.cond(comp, lambda x_, a: torch.einsum('j,kj->k', x_, a), lambda x_, a: torch.einsum('j,jk->k', x_, a),
-                   (x, A,))
+    x = torch.einsum('j,kj->k', x_, a)
     x = x.norm()
     x *= max_abs
     return x

{heavyball-0.20.0.dist-info → heavyball-0.20.1.dist-info}/METADATA

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

{heavyball-0.20.0.dist-info → heavyball-0.20.1.dist-info}/RECORD

@@ -16,9 +16,9 @@ heavyball/precond_schedule_sfpsoap.py,sha256=PNneiOkrRyV1yIZn91lPmYofd1_OiLqJTDy
 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=ESazD0yv14Aa8XKi_pz2CyfVkpcbgYcG2-WMvhQOnxk,35719
-heavyball-0.20.0.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.20.0.dist-info/METADATA,sha256=dJ43LOTrNqh7cDTDzZDSu57goP1gNhU3dfZ26BUK9hA,11926
-heavyball-0.20.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.20.0.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.20.0.dist-info/RECORD,,
+heavyball/utils.py,sha256=14vt4r_MeTsp1q3m0lpgF-Q3PCJg6GLGJrhjRxnbWwQ,35174
+heavyball-0.20.1.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.20.1.dist-info/METADATA,sha256=qzF2P7e2EREeTy_4h85tvUY53omjNm32z83CUHTqt3U,11926
+heavyball-0.20.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.20.1.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.20.1.dist-info/RECORD,,