torchzero 0.4.1__py3-none-any.whl → 0.4.2__py3-none-any.whl

torchzero/modules/basis/soap_basis.py

@@ -0,0 +1,254 @@
+from operator import itemgetter
+import warnings
+
+import torch
+
+from ...core import TensorTransform, Chainable, Module
+from ..adaptive import Adam
+from ...utils import unpack_dicts, unpack_states, TensorList, NumberList, set_storage_
+from ...modules.adaptive.shampoo import _merge_small_dims, _unmerge_small_dims
+from ...linalg import torch_linalg
+from ..adaptive.soap import get_orthogonal_matrix, project, project_back, update_soap_covariances_
+
+# function from https://github.com/nikhilvyas/SOAP/blob/main/soap.py#L240
+@torch.no_grad
+def get_orthogonal_matrix_QR(grad_sqs: list[torch.Tensor], GG: list[torch.Tensor | None], Q_list: list[torch.Tensor | None]):
+    """
+    Computes the eigenbases of the preconditioner using one round of power iteration
+    followed by torch.linalg.qr decomposition.
+     """
+    final = []
+
+    for ind, (M, O) in enumerate(zip(GG, Q_list)):
+
+        # skip 1d or large dims
+        if M is None:
+            final.append(None)
+            continue
+
+        assert O is not None
+
+        est_eig = torch.diagonal(O.T @ M @ O)
+        sort_idx = torch.argsort(est_eig, descending=True)
+        grad_sqs = [s.index_select(ind, sort_idx) for s in grad_sqs]
+
+        power_iter = M @ O[:, sort_idx]
+        Q, _ = torch_linalg.qr(power_iter.to(torch.float32), retry_float64=True)
+        Q = Q.to(power_iter.dtype)
+
+        final.append(Q)
+
+    return final, grad_sqs
+
+class SOAPBasis(TensorTransform):
+    """
+    Run another optimizer in Shampoo eigenbases.
+
+    Note:
+        the buffers of the ``basis_opt`` are re-projected whenever basis changes. The reprojection logic is not implemented on all modules. Some supported modules are:
+
+        ``Adagrad``, ``Adam``, ``Adan``, ``Lion``, ``MARSCorrection``, ``MSAMMomentum``, ``RMSprop``, ``EMA``, ``HeavyBall``, ``NAG``, ``ClipNormByEMA``, ``ClipValueByEMA``, ``NormalizeByEMA``, ``ClipValueGrowth``, ``CoordinateMomentum``, ``CubicAdam``.
+
+        Additionally most modules with no internal buffers are supported, e.g. ``Cautious``, ``Sign``, ``ClipNorm``, ``Orthogonalize``, etc. However modules that use weight values, such as ``WeighDecay`` can't be supported, as weights can't be projected.
+
+        Also, if you say use ``EMA`` on output of ``Pow(2)``, the exponential average will be reprojected as gradient and not as squared gradients. Use modules like ``EMASquared``, ``SqrtEMASquared`` to get correct reprojections.
+
+    Args:
+        basis_opt (Chainable): module or modules to run in Shampoo eigenbases.
+        shampoo_beta (float | None, optional):
+            beta for covariance matrices accumulators. Can be None, then it just sums them like Adagrad (which works worse). Defaults to 0.95.
+        precond_freq (int, optional): How often to update the preconditioner. Defaults to 10.
+        merge_small (bool, optional): Whether to merge small dims. Defaults to True.
+        max_dim (int, optional): Won't precondition dims larger than this. Defaults to 10_000.
+        precondition_1d (bool, optional):
+            Whether to precondition 1d params (SOAP paper sets this to False). Defaults to True.
+        inner (Chainable | None, optional):
+            output of this module is projected and ``basis_opt`` will run on it, but preconditioners are updated
+            from original gradients.
+
+    Examples:
+    SOAP with MARS and AMSGrad:
+    ```python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.SOAPBasis([tz.m.MARSCorrection(0.95), tz.m.Adam(0.95, 0.95, amsgrad=True)]),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    LaProp in Shampoo eigenbases (SOLP):
+    ```python
+
+    # we define LaProp through other modules, moved it out for brevity
+    laprop = (
+        tz.m.RMSprop(0.95),
+        tz.m.Debias(beta1=None, beta2=0.95),
+        tz.m.EMA(0.95),
+        tz.m.Debias(beta1=0.95, beta2=None),
+    )
+
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.SOAPBasis(laprop),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    Lion in Shampoo eigenbases (works kinda well):
+    ```python
+
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.SOAPBasis(tz.m.Lion()),
+        tz.m.LR(1e-3)
+    )
+    ```
+    """
+    def __init__(
+        self,
+        basis_opt: Chainable,
+        shampoo_beta: float | None = 0.95,
+        precond_freq: int = 10,
+        merge_small: bool = True,
+        max_dim: int = 4096,
+        precondition_1d: bool = True,
+        inner: Chainable | None = None,
+    ):
+        defaults = locals().copy()
+        del defaults['self'], defaults["inner"], defaults["basis_opt"]
+
+        super().__init__(defaults)
+        self.set_child("inner", inner)
+        self.set_child("basis_opt", basis_opt)
+
+    @torch.no_grad
+    def single_tensor_initialize(self, tensor, param, grad, loss, state, setting):
+        if setting["merge_small"]:
+            tensor, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(tensor, setting["max_dim"])
+
+        state["exp_avg_proj"] = torch.zeros_like(tensor)
+        state["exp_avg_sq_proj"] = torch.zeros_like(tensor)
+
+        if tensor.ndim <= 1 and not setting["precondition_1d"]:
+            state['GG'] = []
+
+        else:
+            max_dim = setting["max_dim"]
+            state['GG'] = [
+                torch.zeros(s, s, dtype=tensor.dtype, device=tensor.device) if 1<s<max_dim else None for s in tensor.shape
+            ]
+
+        # either scalar parameter, 1d with precondition_1d=False, or all dims are too big.
+        if len([i is not None for i in state['GG']]) == 0:
+            state['GG'] = None
+
+        # first covariance accumulation
+        if state['GG'] is not None:
+            update_soap_covariances_(tensor, GGs_=state['GG'], beta=setting["shampoo_beta"])
+
+            # get projection matrix with first gradients with eigh
+            try: state['Q'] = get_orthogonal_matrix(state['GG'])
+            except torch.linalg.LinAlgError as e:
+                warnings.warn(f"torch.linalg.eigh raised an error when initializing SOAP Q matrices on 1st step, diagonal preconditioning will be used for this parameter. The error was:\n{e}")
+                state["GG"] = None
+
+        state['step'] = 0
+
+
+    # no update to avoid running merge_dims twice
+
+    @torch.no_grad
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
+        # note
+        # do not modify tensors in-place
+        # because they are used to update preconditioner at the end
+
+        steps = [s["step"] for s in states]
+        if any(s == 0 for s in steps):
+            # skip 1st update so to avoid using current gradient in the projection
+            # I scale it instead to avoid issues with further modules
+            for s in states: s["step"] += 1
+            return TensorList(tensors).clamp(-0.1, 0.1)
+            # return TensorList(tensors).zero_()
+
+        merged_updates = [] # for when exp_avg is maintained unprojected
+        merged_grads = [] # this doesn't go into preconditioner
+        projected = []
+
+        # -------------------------------- inner step -------------------------------- #
+        updates = tensors
+        has_inner = "inner" in self.children
+        if has_inner:
+            updates = self.inner_step_tensors("inner", updates, clone=True,
+                                              params=params, grads=grads, loss=loss)
+
+        # ---------------------------------- project --------------------------------- #
+        for grad, update, state, setting in zip(tensors, updates, states, settings):
+            if setting["merge_small"]:
+                update, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(update, setting["max_dim"])
+                if has_inner: # grad is a different tensor, merge it too
+                    grad, _, _ = _merge_small_dims(grad, setting["max_dim"])
+                else: # in this case update is still just grad
+                    grad = update
+
+            merged_updates.append(update)
+            merged_grads.append(grad)
+
+            if state['GG'] is not None:
+                update = project(update, state['Q'])
+
+            projected.append(update)
+
+
+        # ------------------------ run opt in projected space ----------------------- #
+        dirs_proj = self.inner_step_tensors("basis_opt", tensors=projected, clone=True, grads=projected)
+
+        # ------------------------------- project back ------------------------------- #
+        dirs: list[torch.Tensor] = []
+        for dir, state, setting in zip(dirs_proj, states, settings):
+            if state['GG'] is not None:
+                dir = project_back(dir, state['Q'])
+
+            if setting["merge_small"]:
+                dir = _unmerge_small_dims(dir, state['flat_sizes'], state['sort_idxs'])
+
+            dirs.append(dir)
+
+        # -------------------------- update preconditioners -------------------------- #
+        # Update is done after the gradient step to avoid using current gradients in the projection.
+
+        grad_buffs = self.get_child_projected_buffers("basis_opt", "grad")
+        grad_sq_buffs = self.get_child_projected_buffers("basis_opt", ["grad_sq", "grad_cu"])
+
+        for i, (grad, state, setting) in enumerate(zip(merged_grads, states, settings)):
+            if state['GG'] is not None:
+
+                # lerp covariances
+                update_soap_covariances_(grad, state['GG'], beta=setting["shampoo_beta"])
+
+                # (state['step'] - 1) since we start updating on 2nd step
+                if (state['step'] - 1) % setting['precond_freq'] == 0:
+                    g_buffs = [b[i] for b in grad_buffs]
+                    g_sq_buffs = [b[i] for b in grad_sq_buffs]
+
+                    # unproject grad buffers before updating
+                    g_buffs_unproj = [project_back(buff, state["Q"]) for buff in g_buffs]
+
+                    # update projection matrix and exp_avg_sq_proj
+                    try:
+                        state['Q'], g_sq_buffs_new = get_orthogonal_matrix_QR(
+                            g_sq_buffs, state['GG'], state['Q'])
+
+                        for b_old, b_new in zip(g_sq_buffs, g_sq_buffs_new):
+                            set_storage_(b_old, b_new)
+
+                        # re-project grad buffers
+                        for b_proj, b_unproj in zip(g_buffs, g_buffs_unproj):
+                            set_storage_(b_proj, project(b_unproj, state["Q"]))
+
+                    except torch.linalg.LinAlgError:
+                        pass
+
+            state["step"] += 1
+
+        return dirs

torchzero/modules/clipping/ema_clipping.py

@@ -27,60 +27,67 @@ class ClipNormByEMA(TensorTransform):
         self,
         beta=0.99,
         ord: Metrics = 2,
-        eps=1e-6,
         tensorwise:bool=True,
         max_ema_growth: float | None = 1.5,
-        ema_init: Literal['zeros', 'update'] = 'zeros',
+        init: float = 0.0,
+        min_norm: float = 1e-6,
+
         inner: Chainable | None = None,
     ):
-        defaults = dict(beta=beta, ord=ord, tensorwise=tensorwise, ema_init=ema_init, eps=eps, max_ema_growth=max_ema_growth)
+        defaults = dict(beta=beta, ord=ord, tensorwise=tensorwise, init=init, min_norm=min_norm, max_ema_growth=max_ema_growth)
         super().__init__(defaults, inner=inner)
+        self.add_projected_keys("grad", "exp_avg")
 
     @torch.no_grad
     def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
         tensors = TensorList(tensors)
-        ord, tensorwise, ema_init, max_ema_growth = itemgetter('ord', 'tensorwise', 'ema_init', 'max_ema_growth')(settings[0])
+        eps = torch.finfo(tensors[0].dtype).tiny * 2
+        ord, tensorwise, init, max_ema_growth = itemgetter('ord', 'tensorwise', 'init', 'max_ema_growth')(settings[0])
 
-        beta, eps = unpack_dicts(settings, 'beta', 'eps', cls=NumberList)
+        beta, min_norm = unpack_dicts(settings, 'beta', 'min_norm', cls=NumberList)
 
-        ema = unpack_states(states, tensors, 'ema', init = (torch.zeros_like if ema_init=='zeros' else tensors), cls=TensorList)
+        exp_avg = unpack_states(states, tensors, 'exp_avg', init = lambda x: torch.full_like(x, init), cls=TensorList)
 
-        ema.lerp_(tensors, 1-beta)
+        exp_avg.lerp_(tensors, 1-beta)
 
+        # ----------------------------- tensorwise update ---------------------------- #
         if tensorwise:
-            ema_norm = ema.metric(ord)
+            tensors_norm = tensors.norm(ord)
+            ema_norm = exp_avg.metric(ord)
 
             # clip ema norm growth
             if max_ema_growth is not None:
                 prev_ema_norm = unpack_states(states, tensors, 'prev_ema_norm', init=ema_norm, cls=TensorList)
-                allowed_norm = (prev_ema_norm * max_ema_growth).clip(min=1e-6)
+                allowed_norm = (prev_ema_norm * max_ema_growth).clip(min=min_norm)
+
                 ema_denom = (ema_norm / allowed_norm).clip(min=1)
-                ema.div_(ema_denom)
+                exp_avg.div_(ema_denom)
                 ema_norm.div_(ema_denom)
+
                 prev_ema_norm.set_(ema_norm)
 
-            tensors_norm = tensors.norm(ord)
-            denom = tensors_norm / ema_norm.clip(min=eps)
-            if self.NORMALIZE: denom.clip_(min=eps)
-            else: denom.clip_(min=1)
 
+        # ------------------------------- global update ------------------------------ #
         else:
-            ema_norm = ema.global_metric(ord)
+            tensors_norm = tensors.global_metric(ord)
+            ema_norm = exp_avg.global_metric(ord)
 
             # clip ema norm growth
             if max_ema_growth is not None:
                 prev_ema_norm = self.global_state.setdefault('prev_ema_norm', ema_norm)
-                allowed_norm = prev_ema_norm * max_ema_growth
+                allowed_norm = (prev_ema_norm * max_ema_growth).clip(min=min_norm[0])
+
                 if ema_norm > allowed_norm:
-                    ema.div_(ema_norm / allowed_norm)
+                    exp_avg.div_(ema_norm / allowed_norm)
                     ema_norm = allowed_norm
+
                 prev_ema_norm.set_(ema_norm)
 
-            tensors_norm = tensors.global_metric(ord)
-            denom = tensors_norm / ema_norm.clip(min=eps[0])
-            if self.NORMALIZE: denom.clip_(min=eps[0])
-            else: denom.clip_(min=1)
 
+        # ------------------- compute denominator to clip/normalize ------------------ #
+        denom = tensors_norm / ema_norm.clip(min=eps)
+        if self.NORMALIZE: denom.clip_(min=eps)
+        else: denom.clip_(min=1)
         self.global_state['denom'] = denom
 
     @torch.no_grad
@@ -121,7 +128,7 @@ class ClipValueByEMA(TensorTransform):
     def __init__(
         self,
         beta=0.99,
-        init: Literal['zeros', 'update'] = 'zeros',
+        init: float = 0,
 
         inner: Chainable | None = None,
         exp_avg_tfm:Chainable | None=None,
@@ -130,12 +137,10 @@ class ClipValueByEMA(TensorTransform):
         super().__init__(defaults, inner=inner)
 
         self.set_child('exp_avg', exp_avg_tfm)
+        self.add_projected_keys("grad", "exp_avg")
 
     def single_tensor_initialize(self, tensor, param, grad, loss, state, setting):
-        if setting["init"] == "zeros":
-            state["exp_avg"] = torch.zeros_like(tensor)
-        else:
-            state["exp_avg"] = tensor.abs()
+        state["exp_avg"] = tensor.abs() * setting["init"]
 
     @torch.no_grad
     def multi_tensor_update(self, tensors, params, grads, loss, states, settings):
@@ -153,4 +158,4 @@ class ClipValueByEMA(TensorTransform):
             self.inner_step_tensors("exp_avg", exp_avg, clone=True, params=params, grads=grads, loss=loss, must_exist=False))
 
         tensors.clip_(-exp_avg, exp_avg)
-        return tensors
+        return tensors

torchzero/modules/clipping/growth_clipping.py

@@ -30,6 +30,7 @@ class ClipValueGrowth(TensorTransform):
     ):
         defaults = dict(add=add, mul=mul, min_value=min_value, max_decay=max_decay)
         super().__init__(defaults)
+        self.add_projected_keys("grad", "prev")
 
 
     def single_tensor_apply(self, tensor, param, grad, loss, state, setting):

torchzero/modules/experimental/__init__.py

@@ -1,13 +1,9 @@
-"""Those are various ideas of mine plus some other modules that I decided not to move to other sub-packages for whatever reason. This is generally less tested and shouldn't be used."""
-from .adanystrom import AdaNystrom
-from .common_directions_whiten import CommonDirectionsWhiten
+"""Those are various ideas of mine plus some other modules that I decided not to move to other sub-packages for whatever reason. This is generally less tested."""
 from .coordinate_momentum import CoordinateMomentum
 from .cubic_adam import CubicAdam, SubspaceCubicAdam
 from .curveball import CurveBall
-from .eigen_sr1 import EigenSR1
 
 # from dct import DCTProjection
-from .eigengrad import Eigengrad
 from .fft import FFTProjection
 from .gradmin import GradMin
 from .higher_order_newton import HigherOrderNewton
@@ -16,5 +12,4 @@ from .newton_solver import NewtonSolver
 from .newtonnewton import NewtonNewton
 from .reduce_outward_lr import ReduceOutwardLR
 from .scipy_newton_cg import ScipyNewtonCG
-from .spsa1 import SPSA1
 from .structural_projections import BlockPartition, TensorizeProjection

torchzero/modules/experimental/coordinate_momentum.py

@@ -29,6 +29,8 @@ class CoordinateMomentum(TensorTransform):
         defaults = dict(p=p)
         super().__init__(defaults)
 
+        self.add_projected_keys("grad", "velocity")
+
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         p = NumberList(s['p'] for s in settings)

torchzero/modules/experimental/cubic_adam.py

@@ -88,6 +88,10 @@ class CubicAdam(TensorTransform):
         defaults=dict(beta1=beta1,beta2=beta2,beta3=beta3,eps=eps,debiased=debiased,alpha=alpha,mode=mode)
         super().__init__(defaults)
 
+        self.add_projected_keys("grad", "exp_avg")
+        self.add_projected_keys("grad_sq", "exp_avg_sq")
+        self.add_projected_keys("grad_cu", "exp_avg_cu")
+
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1

torchzero/modules/grad_approximation/__init__.py

@@ -1,4 +1,5 @@
-from .grad_approximator import GradApproximator, GradTarget
 from .fdm import FDM
-from .rfdm import RandomizedFDM, MeZO, SPSA, RDSA, GaussianSmoothing
 from .forward_gradient import ForwardGradient
+from .grad_approximator import GradApproximator, GradTarget
+from .rfdm import RDSA, SPSA, GaussianSmoothing, MeZO, RandomizedFDM
+from .spsa1 import SPSA1

torchzero/modules/least_squares/gn.py

@@ -1,3 +1,5 @@
+import warnings
+
 import torch
 
 from ...core import Chainable, Transform
@@ -129,6 +131,10 @@ class GaussNewton(Transform):
                 r = objective.get_loss(backward=False) # n_residuals
                 assert isinstance(r, torch.Tensor)
 
+        if r.numel() == 1:
+            r = r.view(1,1)
+            warnings.warn("Gauss-newton got a single residual. Make sure objective function returns a vector of residuals.")
+
         # set sum of squares scalar loss and it's gradient to objective
         objective.loss = r.pow(2).sum()
 

torchzero/modules/misc/gradient_accumulation.py

@@ -35,6 +35,7 @@ class GradientAccumulation(Module):
     def __init__(self, n: int, mean=True, stop=True):
         defaults = dict(n=n, mean=mean, stop=stop)
         super().__init__(defaults)
+        self.add_projected_keys("grad", "accumulator")
 
 
     @torch.no_grad

torchzero/modules/misc/misc.py

@@ -25,6 +25,7 @@ class Previous(TensorTransform):
         defaults = dict(n=n)
         super().__init__(defaults=defaults)
 
+        self.add_projected_keys("grad", "history")
 
     @torch.no_grad
     def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
@@ -42,6 +43,7 @@ class LastDifference(TensorTransform):
     """Outputs difference between past two updates."""
     def __init__(self,):
         super().__init__()
+        self.add_projected_keys("grad", "prev_tensors")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -54,6 +56,7 @@ class LastGradDifference(Module):
     """Outputs difference between past two gradients."""
     def __init__(self):
         super().__init__()
+        self.add_projected_keys("grad", "prev_grad")
 
     @torch.no_grad
     def apply(self, objective):
@@ -84,6 +87,7 @@ class LastProduct(TensorTransform):
     """Outputs difference between past two updates."""
     def __init__(self):
         super().__init__()
+        self.add_projected_keys("grad", "prev")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -97,6 +101,7 @@ class LastRatio(TensorTransform):
     def __init__(self, numerator: Literal['cur', 'prev'] = 'cur'):
         defaults = dict(numerator=numerator)
         super().__init__(defaults)
+        self.add_projected_keys("grad", "prev")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -112,6 +117,7 @@ class LastAbsoluteRatio(TensorTransform):
     def __init__(self, numerator: Literal['cur', 'prev'] = 'cur', eps:float=1e-8):
         defaults = dict(numerator=numerator, eps=eps)
         super().__init__(defaults)
+        self.add_projected_keys("grad", "prev")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):

torchzero/modules/momentum/averaging.py

@@ -20,6 +20,8 @@ class Averaging(TensorTransform):
         defaults = dict(history_size=history_size)
         super().__init__(defaults=defaults)
 
+        self.add_projected_keys("grad", "history", "average")
+
     @torch.no_grad
     def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
         history_size = setting['history_size']
@@ -45,6 +47,8 @@ class WeightedAveraging(TensorTransform):
         defaults = dict(weights = tolist(weights))
         super().__init__(defaults=defaults)
 
+        self.add_projected_keys("grad", "history")
+
     @torch.no_grad
     def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
         weights = setting['weights']
@@ -79,6 +83,8 @@ class MedianAveraging(TensorTransform):
         defaults = dict(history_size = history_size)
         super().__init__(defaults=defaults)
 
+        self.add_projected_keys("grad", "history")
+
     @torch.no_grad
     def single_tensor_apply(self, tensor, param, grad, loss, state, setting):
         history_size = setting['history_size']

torchzero/modules/momentum/momentum.py

@@ -24,6 +24,8 @@ class EMA(TensorTransform):
         defaults = dict(momentum=momentum,dampening=dampening,debiased=debiased,lerp=lerp,ema_init=ema_init)
         super().__init__(defaults, uses_grad=False)
 
+        self.add_projected_keys("grad", "exp_avg")
+
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
@@ -88,6 +90,8 @@ class NAG(TensorTransform):
         defaults = dict(momentum=momentum,dampening=dampening, lerp=lerp)
         super().__init__(defaults, uses_grad=False)
 
+        self.add_projected_keys("grad", "velocity")
+
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)

torchzero/modules/ops/__init__.py

@@ -13,7 +13,6 @@ from .binary import (
     CopySign,
     Div,
     GraftInputToOutput,
-    GraftInputToOutput,
     GramSchimdt,
     Maximum,
     Minimum,

torchzero/modules/ops/accumulate.py

@@ -13,6 +13,7 @@ class AccumulateSum(TensorTransform):
     def __init__(self, decay: float = 0):
         defaults = dict(decay=decay)
         super().__init__(defaults)
+        self.add_projected_keys("grad", "sum")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -30,6 +31,7 @@ class AccumulateMean(TensorTransform):
     def __init__(self, decay: float = 0):
         defaults = dict(decay=decay)
         super().__init__(defaults)
+        self.add_projected_keys("grad", "mean")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -65,6 +67,7 @@ class AccumulateMaximum(TensorTransform):
     def __init__(self, decay: float = 0):
         defaults = dict(decay=decay)
         super().__init__(defaults)
+        self.add_projected_keys("grad", "maximum")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -82,6 +85,7 @@ class AccumulateMinimum(TensorTransform):
     def __init__(self, decay: float = 0):
         defaults = dict(decay=decay)
         super().__init__(defaults)
+        self.add_projected_keys("grad", "minimum")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):

torchzero/modules/ops/higher_level.py

@@ -30,6 +30,7 @@ class EMASquared(TensorTransform):
     def __init__(self, beta:float=0.999, amsgrad=False, pow:float=2):
         defaults = dict(beta=beta,pow=pow,amsgrad=amsgrad)
         super().__init__(defaults)
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "max_exp_avg_sq")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -57,7 +58,7 @@ class SqrtEMASquared(TensorTransform):
     def __init__(self, beta:float=0.999, amsgrad=False, debiased: bool = False, pow:float=2,):
         defaults = dict(beta=beta,pow=pow,amsgrad=amsgrad,debiased=debiased)
         super().__init__(defaults)
-
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "max_exp_avg_sq")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -141,6 +142,8 @@ class CenteredEMASquared(TensorTransform):
     def __init__(self, beta: float = 0.99, amsgrad=False, pow:float=2):
         defaults = dict(beta=beta, amsgrad=amsgrad, pow=pow)
         super().__init__(defaults, uses_grad=False)
+        self.add_projected_keys("grad", "exp_avg")
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "max_exp_avg_sq")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
@@ -175,6 +178,8 @@ class CenteredSqrtEMASquared(TensorTransform):
     def __init__(self, beta: float = 0.99, amsgrad=False, debiased: bool = False, pow:float=2):
         defaults = dict(beta=beta, amsgrad=amsgrad, debiased=debiased, pow=pow)
         super().__init__(defaults, uses_grad=False)
+        self.add_projected_keys("grad", "exp_avg")
+        self.add_projected_keys("grad_sq", "exp_avg_sq", "max_exp_avg_sq")
 
     @torch.no_grad
     def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):

torchzero/modules/second_order/inm.py

@@ -35,6 +35,8 @@ class ImprovedNewton(Transform):
         self,
         damping: float = 0,
         eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        eigv_tol: float | None = None,
+        truncate: int | None = None,
         update_freq: int = 1,
         precompute_inverse: bool | None = None,
         use_lstsq: bool = False,
@@ -89,6 +91,8 @@ class ImprovedNewton(Transform):
             state = self.global_state,
             damping = fs["damping"],
             eigval_fn = fs["eigval_fn"],
+            eigv_tol = fs["eigv_tol"],
+            truncate = fs["truncate"],
             precompute_inverse = precompute_inverse,
             use_lstsq = fs["use_lstsq"]
         )