torchzero 0.4.2__tar.gz → 0.4.3__tar.gz

{torchzero-0.4.2 → torchzero-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.4.2
+Version: 0.4.3
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 Project-URL: Homepage, https://github.com/inikishev/torchzero

{torchzero-0.4.2 → torchzero-0.4.3}/pyproject.toml

@@ -13,7 +13,7 @@ build-backend = "setuptools.build_meta"
 name = "torchzero"
 description = "Modular optimization library for PyTorch."
 
-version = "0.4.2"
+version = "0.4.3"
 dependencies = [
   "torch",
   "numpy",

{torchzero-0.4.2 → torchzero-0.4.3}/tests/test_identical.py

@@ -105,7 +105,7 @@ def test_adam(amsgrad):
     tz_fn_ops = lambda p: tz.Optimizer(
         p,
         tz.m.DivModules(
-            tz.m.EMA(0.9, debiased=True),
+            tz.m.EMA(0.9, debias=True),
             [tz.m.SqrtEMASquared(0.999, debiased=True, amsgrad=amsgrad), tz.m.Add(1e-8)]
         ))
     tz_fn_ops2 = lambda p: tz.Optimizer(

{torchzero-0.4.2 → torchzero-0.4.3}/torchzero/modules/basis/ggt_basis.py

@@ -111,7 +111,7 @@ class GGTBasis(TensorTransform):
         inner: Chainable | None = None,
     ):
         defaults = locals().copy()
-        del defaults['self'], defaults['inner']
+        del defaults['self'], defaults['inner'], defaults["basis_opt"]
 
         super().__init__(defaults, concat_params=True, inner=inner)
         self.set_child("basis_opt", basis_opt)

{torchzero-0.4.2 → torchzero-0.4.3}/torchzero/modules/momentum/momentum.py

@@ -6,7 +6,7 @@ import torch
 
 from ...core import  TensorTransform
 from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
-from ..opt_utils import debias, ema_
+from ..opt_utils import debias as _debias, ema_
 
 
 class EMA(TensorTransform):
@@ -15,13 +15,13 @@ class EMA(TensorTransform):
     Args:
         momentum (float, optional): momentum (beta). Defaults to 0.9.
         dampening (float, optional): momentum dampening. Defaults to 0.
-        debiased (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
+        debias (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
         lerp (bool, optional): whether to use linear interpolation. Defaults to True.
         ema_init (str, optional): initial values for the EMA, "zeros" or "update".
         target (Target, optional): target to apply EMA to. Defaults to 'update'.
     """
-    def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=True, ema_init: Literal['zeros', 'update'] = 'zeros'):
-        defaults = dict(momentum=momentum,dampening=dampening,debiased=debiased,lerp=lerp,ema_init=ema_init)
+    def __init__(self, momentum:float=0.9, dampening:float=0, debias: bool = False, lerp=True, ema_init: Literal['zeros', 'update'] = 'zeros'):
+        defaults = dict(momentum=momentum,dampening=dampening,debias=debias,lerp=lerp,ema_init=ema_init)
         super().__init__(defaults, uses_grad=False)
 
         self.add_projected_keys("grad", "exp_avg")
@@ -30,7 +30,7 @@ class EMA(TensorTransform):
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
-        debiased, lerp, ema_init = itemgetter('debiased','lerp','ema_init')(settings[0])
+        debias, lerp, ema_init = itemgetter('debias','lerp','ema_init')(settings[0])
 
         exp_avg = unpack_states(states, tensors, 'exp_avg',
                                 init=torch.zeros_like if ema_init=='zeros' else tensors, cls=TensorList)
@@ -38,7 +38,7 @@ class EMA(TensorTransform):
 
         exp_avg = ema_(TensorList(tensors), exp_avg_=exp_avg,beta=momentum,dampening=dampening,lerp=lerp)
 
-        if debiased: return debias(exp_avg, step=step, beta1=momentum, alpha=1, inplace=False)
+        if debias: return _debias(exp_avg, step=step, beta1=momentum, alpha=1, inplace=False)
         else: return exp_avg.clone() # this has exp_avg storage so needs to be cloned
 
 
@@ -49,14 +49,14 @@ class HeavyBall(EMA):
     Args:
         momentum (float, optional): momentum (beta). Defaults to 0.9.
         dampening (float, optional): momentum dampening. Defaults to 0.
-        debiased (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
+        debias (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
         lerp (bool, optional):
             whether to use linear interpolation, if True, this becomes exponential moving average. Defaults to False.
         ema_init (str, optional): initial values for the EMA, "zeros" or "update".
         target (Target, optional): target to apply EMA to. Defaults to 'update'.
     """
-    def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=False, ema_init: Literal['zeros', 'update'] = 'update'):
-        super().__init__(momentum=momentum, dampening=dampening, debiased=debiased, lerp=lerp, ema_init=ema_init)
+    def __init__(self, momentum:float=0.9, dampening:float=0, debias: bool = False, lerp=False, ema_init: Literal['zeros', 'update'] = 'update'):
+        super().__init__(momentum=momentum, dampening=dampening, debias=debias, lerp=lerp, ema_init=ema_init)
 
 def nag_(
     tensors_: TensorList,

{torchzero-0.4.2 → torchzero-0.4.3}/torchzero/utils/optuna_tools.py

@@ -27,7 +27,7 @@ def get_momentum(trial: optuna.Trial, prefix: str, conditional: bool=True) -> li
             m = NAG(beta, dampening, lerp)
             if debiased: m = Chain(m, Debias(beta1=beta))
         else:
-            m = EMA(beta, dampening, debiased=debiased, lerp=lerp)
+            m = EMA(beta, dampening, debias=debiased, lerp=lerp)
         return [m]
     return []
 

{torchzero-0.4.2 → torchzero-0.4.3}/torchzero.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.4.2
+Version: 0.4.3
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 Project-URL: Homepage, https://github.com/inikishev/torchzero