adv-optm 1.2.dev5__tar.gz → 1.2.dev6__tar.gz

{adv_optm-1.2.dev5 → adv_optm-1.2.dev6}/PKG-INFO

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

{adv_optm-1.2.dev5 → adv_optm-1.2.dev6}/adv_optm/__init__.py

@@ -20,4 +20,4 @@ __all__ = [
     "AdaMuon_adv",
 ]
 
-__version__ = "1.2.dev5"
+__version__ = "1.2.dev6"

{adv_optm-1.2.dev5 → adv_optm-1.2.dev6}/adv_optm/util/Kourkoutas.py

@@ -86,9 +86,17 @@ class KourkoutasHelper:
         # These are just for the sample log, initialize them
         sun, pooled_grad_norm, prev_r_ema_val, r_ema_tensor = (torch.tensor(0.0),)*4
 
+        # The optimizer that owns this helper holds the master defaults for K-b.
+        # This is crucial in hybrid optimizers where some param_groups might not
+        # have all K-b keys populated, preventing KeyErrors.
+        master_defaults = self.optimizer.defaults
+
         for layer_key, info in self.layer_info.items():
             params, group = info['params'], info['group_ref']
 
+            if not group.get('kourkoutas_beta', False):
+                continue
+
             first_param_in_layer = info['params'][0]
             param_state = self.optimizer.state[first_param_in_layer]
 
@@ -100,6 +108,15 @@ class KourkoutasHelper:
             if 'kourkoutas_r_ema' not in param_state:
                 param_state['kourkoutas_r_ema'] = torch.tensor(0.0, device=first_param_in_layer.device, dtype=torch.float32)
 
+            # Use group-specific K-b settings, falling back to the optimizer's master defaults.
+            # This makes the helper robust against param groups that enable kourkoutas_beta
+            # but are missing the other required hyperparameters.
+            ema_alpha = group.get('ema_alpha', master_defaults['ema_alpha'])
+            beta2_max = group.get('betas', master_defaults['betas'])[1]
+            beta2_min = group.get('beta2_min', master_defaults['beta2_min'])
+            tiny_spike = group.get('tiny_spike', master_defaults['tiny_spike'])
+            k_warmup_steps = group.get('k_warmup_steps', master_defaults['k_warmup_steps'])
+
             r_ema_tensor = param_state['kourkoutas_r_ema']
             accumulator = self.layer_state[layer_key]['sum_sq_accumulator']
             
@@ -107,17 +124,16 @@ class KourkoutasHelper:
             prev_r_ema_val = r_ema_tensor.item() # for logging
             
             # Update the persistent EMA tensor in-place.
-            r_ema_tensor.mul_(group['ema_alpha']).add_(pooled_grad_norm, alpha=1.0 - group['ema_alpha'])
+            r_ema_tensor.mul_(ema_alpha).add_(pooled_grad_norm, alpha=1.0 - ema_alpha)
             
-            beta2_max = group['betas'][1]
             sun = torch.tensor(0.0, device=r_ema_tensor.device) # Default sun to 0 for warmup
             
-            if current_step < group['k_warmup_steps']:
+            if current_step < k_warmup_steps:
                 beta2 = beta2_max
             else:
-                raw = pooled_grad_norm / (r_ema_tensor + group['tiny_spike'])
+                raw = pooled_grad_norm / (r_ema_tensor + tiny_spike)
                 sun = raw / (1.0 + raw)
-                beta2 = beta2_max - (beta2_max - group['beta2_min']) * sun
+                beta2 = beta2_max - (beta2_max - beta2_min) * sun
 
             # Store the final calculated beta2 in the helper's transient state for this step.
             self.layer_state[layer_key]['dynamic_beta2'] = beta2.item() if isinstance(beta2, torch.Tensor) else beta2

{adv_optm-1.2.dev5 → adv_optm-1.2.dev6}/adv_optm.egg-info/PKG-INFO

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

{adv_optm-1.2.dev5 → adv_optm-1.2.dev6}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="adv_optm",
-    version="1.2.dev5",
+    version="1.2.dev6",
     author="Koratahiu",
     author_email="hiuhonor@gmail.com",
     license='Apache 2.0',