adv-optm 2.6.2.dev1__py3-none-any.whl → 2.6.2.dev2__py3-none-any.whl

adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "SinkSGD_adv",
 ]
 
-__version__ = "2.6.2.dev1"
+__version__ = "2.6.2.dev2"

adv_optm/optim/SinkSGD_adv.py

@@ -373,9 +373,13 @@ class SinkSGD_adv(torch.optim.Optimizer):
                 del anchor
 
         update_scaling = step_size
-        if group.get('spectral_normalization', False) and not (oft_sym and not normed_mt):
+
+        # If spectral normalization was already applied natively in the OFT sinkhorn, skip the scaling block
+        spectral_sinkhorn_applied = oft_sym and not normed_mt and group.get('spectral_normalization', False)
+
+        if group.get('spectral_normalization', False) and not spectral_sinkhorn_applied:
             update = scale_update(p, update, update_scaling, state=state)
-        else:
+        elif not spectral_sinkhorn_applied:
             if snr_cond:
                 update_scaling = update_scaling * (4/math.pi)
             update.mul_(update_scaling)

adv_optm/util/sinkhorn.py

@@ -174,14 +174,13 @@ def apply_oft_sinkhorn(
             g_orth_norm = torch.linalg.vector_norm(G, ord=2, dim=(1, 2), keepdim=True).clamp_min_(1e-12)
             G.mul_(target_norm / g_orth_norm)
 
-    # Global normalization
-    norm = torch.linalg.vector_norm(G).clamp_min_(1e-12)
-    target_norm_global = math.sqrt(G.numel())
-    G.mul_(target_norm_global / norm)
-
-    # Extract upper triangular elements back into the flat format
     update.copy_(G[batch_idx, rows, cols])
 
+    # Global normalization
+    norm = torch.linalg.vector_norm(update).clamp_min_(1e-12)
+    target_norm_global = math.sqrt(update.numel())
+    update.mul_(target_norm_global / norm)
+
     return update
 
 @torch.no_grad()
@@ -201,7 +200,6 @@ def apply_spectral_oft_sinkhorn(
     device, dtype = p.device, p.dtype
     rows, cols = scaled_optm.get_cached_structural_tensors(block_size, device)
 
-    # Flatten any prepended batch dimensions for processing
     orig_shape = p.shape
 
     # Align the scale of p with the forward pass

{adv_optm-2.6.2.dev1.dist-info → adv_optm-2.6.2.dev2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: adv_optm
-Version: 2.6.2.dev1
+Version: 2.6.2.dev2
 Summary: A family of highly efficient, lightweight yet powerful optimizers.
 Home-page: https://github.com/Koratahiu/Advanced_Optimizers
 Author: Koratahiu

{adv_optm-2.6.2.dev1.dist-info → adv_optm-2.6.2.dev2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-adv_optm/__init__.py,sha256=tbk8Ux10bVPAtT5E52dUr7353gHwov9vu00eicn6glM,357
+adv_optm/__init__.py,sha256=89CN2Q1tBX9qe9GvkxvFPyt5H1QgFRO72ST3UaAGReo,357
 adv_optm/optim/AdaMuon_adv.py,sha256=WX6AbkDYFVqYC5dtAWnOzMoNpysoTA7eSmCBl1eyuQ8,33865
 adv_optm/optim/AdamW_adv.py,sha256=XCJv7jbKyE5iDeag2Gr0BWG0sa91BVze7J-l-3P35vU,23428
 adv_optm/optim/Adopt_adv.py,sha256=nH65TYFG2gPzXWQ6B9tT0dAW0MbF49gSovP2oeMuitQ,23934
@@ -6,7 +6,7 @@ adv_optm/optim/Lion_adv.py,sha256=xdr85vnR0JjDhbAWP3PyH3STnaoObqY6-iys4YR2Gsk,12
 adv_optm/optim/Muon_adv.py,sha256=bIB0yaQ3mTI3bAoAD7hymIshvX8QCdgyYeT61x1dcG4,27970
 adv_optm/optim/Prodigy_adv.py,sha256=AxfPuEZ-tAW9ogNh4R_mVWzmYCnNpT3lfk3bUbjnrq0,27576
 adv_optm/optim/SignSGD_adv.py,sha256=VWLu7YdaouSlwQgl_Q_J5ZrGwohjEeXpleZW1F9KMUI,16860
-adv_optm/optim/SinkSGD_adv.py,sha256=yCGXvT80xpqfXGAqfH-KiF7jKMqvU0DzK9Zo_v1viQ0,17909
+adv_optm/optim/SinkSGD_adv.py,sha256=-OwbZA7MOCxhLaefAKJ_OWSWTn4E9O6E_YaYR28l9Y0,18161
 adv_optm/optim/__init__.py,sha256=RkpzWpEgAOdQAppTBLc5AilfJq-wn0aabNsQDj_7-4A,452
 adv_optm/util/Kourkoutas.py,sha256=Hj7LubFNacZDGDYlFi-9p1yVxi7ACeVSNRhMgQUdjAQ,12283
 adv_optm/util/Muon_AuxAdam.py,sha256=DFTVsnoR4gg7U7d8U18sEc6FB2FWAHUdOh98DCrG7a8,8415
@@ -19,11 +19,11 @@ adv_optm/util/lion_k.py,sha256=b3oqTHK1spoD_Hi8wxKZ9lEWUkHejovT4piU-wUp_eI,1686
 adv_optm/util/param_update.py,sha256=Gj5cwKxP5A3bZU2WNkqTE1GqwEcusRWEMCjxtB82hEo,16149
 adv_optm/util/scaled_optm.py,sha256=HOlsDGqrF7oRiZBG2noJ6JeoxeJK5dOXfQO6P51Mv9I,9864
 adv_optm/util/signed_util.py,sha256=hOmHpHuSdAMZXS21nCeDyoV2prjZ5YJEXNum_l6s0O0,2092
-adv_optm/util/sinkhorn.py,sha256=QWkdIEU8UUHdXMSKRFrcwxtprM0iDYDEMpdlZW1G8pE,11042
+adv_optm/util/sinkhorn.py,sha256=WkEQNE5hA_ktLXc8yRZZqND9Hiw7C1n1ELXKcRF-908,10929
 adv_optm/util/state_util.py,sha256=vVP4smoYNxrsofhFm1YvUixQRI8IsbqElYk3uElmJ2w,10181
 adv_optm/util/update_util.py,sha256=IAaYuxrJ9kP2P4Z0v02pesRe4mRzP5SvXbLHjLrrdZw,1332
-adv_optm-2.6.2.dev1.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
-adv_optm-2.6.2.dev1.dist-info/METADATA,sha256=H4kxfqqFhMDSYN_z30oijEW2Y_HB9oBsMmC8q0y5QG0,5996
-adv_optm-2.6.2.dev1.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
-adv_optm-2.6.2.dev1.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
-adv_optm-2.6.2.dev1.dist-info/RECORD,,
+adv_optm-2.6.2.dev2.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+adv_optm-2.6.2.dev2.dist-info/METADATA,sha256=Xk16b8spXy1nlOLid5JAhPGLB_Y7zU0n0onmm633QoY,5996
+adv_optm-2.6.2.dev2.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+adv_optm-2.6.2.dev2.dist-info/top_level.txt,sha256=iNfBIIzu-lPrQ7jyC56WBCcbkRwitM2nJ15-MRQ_6fg,9
+adv_optm-2.6.2.dev2.dist-info/RECORD,,