torchzero 0.3.10__py3-none-any.whl → 0.3.11__py3-none-any.whl

torchzero/modules/optimizers/msam.py

@@ -0,0 +1,186 @@
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, Target, Transform, apply_transform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states, generic_ne
+from ..functional import ema_
+from ..momentum.momentum import nag_
+
+
+def msam_(
+    tensors: TensorList,
+    params: TensorList,
+    velocity_: TensorList,
+    momentum: float | NumberList,
+    lr: NumberList | None,
+    rho: float | NumberList,
+    weight_decay: float | NumberList,
+    nesterov: bool = False,
+    lerp: bool = False,
+
+    # inner args
+    inner: Module | None = None,
+    grads: list[torch.Tensor] | None = None,
+):
+    # weights w and wh, momentum μ, perturbation strength ρ
+    # w = wh + rho * v / ||v||
+    # v1 = μv + g
+    # w1 = w - lr*v1
+    # wh1 = w1 - rho * v1 / ||v1||
+
+    # w1 = wh + rho * v / ||v|| - lr*v1
+    # vn = rho * v / ||v||
+    # v1n = rho * v1 / ||v1||
+    # wh1 = wh + vn - lr*v1 - v1n
+
+    # the update is
+    # vn - lr*v1 - v1n
+
+    # we track ascent direction so it becomes lr*v1 + v1n - vn
+
+    # can't really decouple it from lr
+    # but at least it is now expressed as function of g
+
+    denom = (velocity_.global_vector_norm() / rho).clip(min=1e-8)
+    vn = velocity_ / denom
+
+    mom_ = nag_ if nesterov else ema_
+    velocity_ = mom_(tensors, velocity_, momentum, dampening=0, lerp=lerp)
+
+    denom = (velocity_.global_vector_norm() / rho).clip(min=1e-8)
+    v1n = velocity_ / denom
+
+    if inner is not None:
+        assert params is not None
+        inner_update = TensorList(apply_transform(inner, tensors, params=params, grads=grads))
+
+    else:
+        assert lr is not None
+        inner_update = velocity_ * lr
+
+    update = inner_update.add_(v1n).sub_(vn)
+
+    if generic_ne(weight_decay, 0):
+        wd = (params + vn).mul_(weight_decay)
+        update.add_(wd)
+
+    return update
+
+class MSAM(Transform):
+    """Momentum-SAM from https://arxiv.org/pdf/2401.12033.
+
+    This implementation expresses the update rule as function of gradient. This way it can be used as a drop-in
+    replacement for momentum strategies in other optimizers.
+
+    To combine MSAM with other optimizers in the way done in the official implementation,
+    e.g. to make Adam_MSAM, use :code:`tz.m.MSAMObjective` module.
+
+    .. note::
+        MSAM has a learning rate hyperparameter that can't really be removed from the update rule.
+        To avoid compounding learning rate mofications, remove the :code:`tz.m.LR` module if you had it.
+
+    Args:
+        lr (float): learning rate. Adding this module adds support for learning rate schedulers.
+        momentum (float, optional): momentum (beta). Defaults to 0.9.
+        rho (float, optional): perturbation strength. Defaults to 0.3.
+        weight_decay (float, optional):
+            weight decay. It is applied to perturbed parameters, so it is differnet
+            from applying :code:`tz.m.WeightDecay` after MSAM. Defaults to 0.
+        nesterov (bool, optional): whether to use nesterov momentum formula. Defaults to False.
+        lerp (bool, optional):
+            whether to use linear interpolation, if True, this becomes similar to exponential moving average. Defaults to False.
+
+    Examples:
+        MSAM
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.MSAM(1e-3)
+            )
+
+        Adam with MSAM instead of exponential average. Note that this is different from Adam_MSAM.
+        To make Adam_MSAM and such, use the :code:`tz.m.MSAMObjective` module.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.RMSprop(0.999, inner=tz.m.MSAM(1e-3)),
+                tz.m.Debias(0.9, 0.999),
+            )
+    """
+    USES_LR = True
+    def __init__(self, lr: float, momentum:float=0.9, rho:float=0.3,  weight_decay:float=0, nesterov=False, lerp=False,):
+        defaults = dict(momentum=momentum,rho=rho, nesterov=nesterov, lerp=lerp, weight_decay=weight_decay)
+        if self.USES_LR: defaults['lr'] = lr
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
+        s = self.settings[params[0]]
+        lerp = s['lerp']
+        nesterov = s['nesterov']
+
+        if self.USES_LR:
+            lr, momentum, rho, weight_decay = unpack_dicts(settings, 'lr','momentum','rho','weight_decay', cls=NumberList)
+
+        else:
+            lr=None
+            momentum,rho,weight_decay = unpack_dicts(settings, 'momentum','rho','weight_decay', cls=NumberList)
+
+        return msam_(
+            TensorList(tensors),
+            params=TensorList(params),
+            velocity_=velocity,
+            momentum=momentum,
+            lr=lr,
+            rho=rho,
+            weight_decay=weight_decay,
+            nesterov=nesterov,
+            lerp=lerp,
+
+            # inner args
+            inner=self.children.get("modules", None),
+            grads=grads,
+        )
+
+
+class MSAMObjective(MSAM):
+    """Momentum-SAM from https://arxiv.org/pdf/2401.12033.
+
+    .. note::
+        Please make sure to place :code:`tz.m.LR` inside the :code:`modules` argument. For example,
+        :code:`tz.m.MSAMObjective([tz.m.Adam(), tz.m.LR(1e-3)])`. Putting LR after MSAM will lead
+        to an incorrect update rule.
+
+    Args:
+        modules (Chainable): modules that will optimizer the MSAM objective. Make sure :code:`tz.m.LR` is one of them.
+        momentum (float, optional): momentum (beta). Defaults to 0.9.
+        rho (float, optional): perturbation strength. Defaults to 0.3.
+        nesterov (bool, optional): whether to use nesterov momentum formula. Defaults to False.
+        lerp (bool, optional):
+            whether to use linear interpolation, if True, MSAM momentum becomes similar to exponential moving average.
+            Defaults to False.
+
+    Examples:
+        AdamW-MSAM
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                bench.parameters(),
+                tz.m.MSAMObjective(
+                    [tz.m.Adam(), tz.m.WeightDecay(1e-3), tz.m.LR(1e-3)],
+                    rho=1.
+                )
+            )
+    """
+    USES_LR = False
+    def __init__(self, modules: Chainable, momentum:float=0.9, rho:float=0.3, weight_decay:float=0, nesterov=False, lerp=False):
+        super().__init__(lr=0, momentum=momentum, rho=rho, weight_decay=weight_decay, nesterov=nesterov, lerp=lerp)
+        self.set_child('modules', modules)
+

torchzero/modules/optimizers/muon.py

@@ -19,6 +19,7 @@ def _is_at_least_2d(p: torch.Tensor):
 
 # stolen from:
 # https://github.com/KellerJordan/Muon/blob/master/muon.py
+# actually at this stage its a frankenstein
 @enable_compilation
 def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int) -> torch.Tensor:
     """
@@ -152,7 +153,7 @@ class Orthogonalize(TensorwiseTransform):
     The Muon page says that embeddings and classifier heads should not be orthogonalized.
     Usually only matrix parameters that are directly used in matmuls should be orthogonalized.
 
-    To make Muon, use Split with Adam on 1d params: TODO code example.
+    To make Muon, use Split with Adam on 1d params
 
     Args:
         ns_steps (int, optional):
@@ -165,6 +166,30 @@ class Orthogonalize(TensorwiseTransform):
             Newton-Schulz is very fast, SVD is extremely slow but can be slighly more precise.
         target (str, optional):
             what to set on var.
+
+
+    Examples:
+        standard Muon with Adam fallback
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.head.parameters(),
+                tz.m.Split(
+                    # apply muon only to 2D+ parameters
+                    filter = lambda t: t.ndim >= 2,
+                    true = [
+                        tz.m.HeavyBall(),
+                        tz.m.Orthogonalize(),
+                        tz.m.LR(1e-2),
+                    ],
+                    false = tz.m.Adam()
+                ),
+                tz.m.LR(1e-2)
+            )
+
+    Reference:
+        Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, Jeremy Bernstein - Muon: An optimizer for hidden layers in neural networks (2024) https://github.com/KellerJordan/Muon
     """
     def __init__(self, ns_steps=5, adjust_lr=False, dual_norm_correction=False,
                  method: Literal['newton-schulz', 'svd'] = 'newton-schulz', target:Target='update'):
@@ -172,9 +197,9 @@ class Orthogonalize(TensorwiseTransform):
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
         orthogonalize, ns_steps, dual_norm_correction, adjust_lr, method = itemgetter(
-            'orthogonalize', 'ns_steps', 'dual_norm_correction', 'adjust_lr', 'method')(settings)
+            'orthogonalize', 'ns_steps', 'dual_norm_correction', 'adjust_lr', 'method')(setting)
 
         if not orthogonalize: return tensor
 
@@ -199,7 +224,7 @@ class DualNormCorrection(TensorwiseTransform):
     def __init__(self, target: Target='update'):
         super().__init__({}, uses_grad=True, target=target)
 
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
         assert grad is not None
         if (tensor.ndim >= 2) and (tensor.size(0) > 1) and (tensor.size(1) > 1):
             return _dual_norm_correction(tensor, grad, batch_first=False)
@@ -213,7 +238,7 @@ class MuonAdjustLR(Transform):
         defaults = dict(alpha=alpha)
         super().__init__(defaults=defaults, uses_grad=False, target=target)
 
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         alphas = [s['alpha'] for s in settings]
         tensors_alphas = [(t, adjust_lr_for_muon(a, t.shape)) for t, a in zip(tensors, alphas) if _is_at_least_2d(t)]
         tensors = [i[0] for i in tensors_alphas]

torchzero/modules/optimizers/orthograd.py

@@ -36,7 +36,7 @@ class OrthoGrad(Transform):
         defaults = dict(eps=eps, renormalize=renormalize)
         super().__init__(defaults, uses_grad=False, target=target)
 
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         eps = settings[0]['eps']
         renormalize = settings[0]['renormalize']
 

torchzero/modules/optimizers/rmsprop.py

@@ -40,7 +40,9 @@ def rmsprop_(
     return tensors_.div_(sqrt_exp_avg_sq.add_(eps))
 
 class RMSprop(Transform):
-    """Divides graient by EMA of gradient squares. Matches pytorch RMSprop if "init" is set to "zeros".
+    """Divides graient by EMA of gradient squares.
+
+    This implementation is identical to :code:`torch.optim.RMSprop`.
 
     Args:
         smoothing (float, optional): beta for exponential moving average of gradient squares. Defaults to 0.99.
@@ -50,7 +52,8 @@ class RMSprop(Transform):
         amsgrad (bool, optional): Whether to divide by maximum of EMA of gradient squares instead. Defaults to False.
         pow (float, optional): power used in second momentum power and root. Defaults to 2.
         init (str, optional): how to initialize EMA, either "update" to use first update or "zeros". Defaults to "update".
-        inner (Chainable | None, optional): Inner modules that are applied after updating EMA and before preconditioning. Defaults to None.
+        inner (Chainable | None, optional):
+            Inner modules that are applied after updating EMA and before preconditioning. Defaults to None.
     """
     def __init__(
         self,
@@ -60,7 +63,7 @@ class RMSprop(Transform):
         debiased: bool = False,
         amsgrad: bool = False,
         pow: float = 2,
-        init: Literal["zeros", "update"] = "update",
+        init: Literal["zeros", "update"] = "zeros",
         inner: Chainable | None = None,
     ):
         defaults = dict(smoothing=smoothing,eps=eps,centered=centered,debiased=debiased,amsgrad=amsgrad,pow=pow,init=init)
@@ -69,7 +72,7 @@ class RMSprop(Transform):
         if inner is not None:
             self.set_child('inner', inner)
 
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
         smoothing, eps = unpack_dicts(settings, 'smoothing', 'eps', cls=NumberList)
         centered, debiased, amsgrad, pow, init = itemgetter('centered','debiased','amsgrad','pow','init')(settings[0])

torchzero/modules/optimizers/rprop.py

@@ -135,7 +135,8 @@ class Rprop(Transform):
     Next step, magnitude for that weight won't change.
 
     Compared to pytorch this also implements backtracking update when sign changes.
-    To make this behave exactly the same as `torch.optim.Rprop`, set `backtrack` to False.
+
+    This implementation is identical to :code:`torch.optim.Rprop` if :code:`backtrack` is set to False.
 
     Args:
         nplus (float): multiplicative increase factor for when ascent didn't change sign (default: 1.2).
@@ -164,7 +165,7 @@ class Rprop(Transform):
         super().__init__(defaults, uses_grad=False)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
@@ -223,7 +224,7 @@ class ScaleLRBySignChange(Transform):
         super().__init__(defaults, uses_grad=use_grad, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
@@ -272,7 +273,7 @@ class BacktrackOnSignChange(Transform):
         super().__init__(defaults, uses_grad=use_grad)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
@@ -294,12 +295,29 @@ class BacktrackOnSignChange(Transform):
         return tensors
 
 class SignConsistencyMask(Transform):
-    """0 if sign changed 1 otherwise"""
+    """
+    Outputs a mask of sign consistency of current and previous inputs.
+
+    The output is 0 for weights where input sign changed compared to previous input, 1 otherwise.
+
+    Examples:
+
+        GD that skips update for weights where gradient sign changed compared to previous gradient.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Mul(tz.m.SignConsistencyMask()),
+                tz.m.LR(1e-2)
+            )
+
+    """
     def __init__(self,target: Target = 'update'):
         super().__init__({}, uses_grad=False, target = target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         prev = unpack_states(states, tensors, 'prev', cls=TensorList)
         mask = prev.mul_(tensors).gt_(0)
         prev.copy_(tensors)
@@ -307,7 +325,23 @@ class SignConsistencyMask(Transform):
 
 
 class SignConsistencyLRs(Transform):
-    """LR for each weight is increased when two consequtive update signs are the same, decreased otherwise. This returns the LRs themselves."""
+    """Outputs per-weight learning rates based on consecutive sign consistency.
+
+    The learning rate for a weight is multiplied by :code:`nplus` when two consecutive update signs are the same, otherwise it is multiplied by :code:`nplus`. The learning rates are bounded to be in :code:`(lb, ub)` range.
+
+    Examples:
+
+        GD scaled by consecutive gradient sign consistency
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Mul(tz.m.SignConsistencyLRs()),
+                tz.m.LR(1e-2)
+            )
+
+    """
     def __init__(
         self,
         nplus: float = 1.2,
@@ -321,7 +355,7 @@ class SignConsistencyLRs(Transform):
         super().__init__(defaults, uses_grad=False, target = target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 

torchzero/modules/optimizers/sam.py

@@ -0,0 +1,163 @@
+from contextlib import nullcontext
+import torch
+from ...utils import TensorList, NumberList
+from ...core import Module
+
+
+class SAM(Module):
+    """Sharpness-Aware Minimization from https://arxiv.org/pdf/2010.01412
+
+    SAM functions by seeking parameters that lie in neighborhoods having uniformly low loss value.
+    It performs two forward and backward passes per step.
+
+    This implementation modifies the closure to return loss and calculate gradients
+    of the SAM objective. All modules after this will use the modified objective.
+
+    .. note::
+        This module requires a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients at two points on each step.
+
+    Args:
+        rho (float, optional): Neighborhood size. Defaults to 0.05.
+        p (float, optional): norm of the SAM objective. Defaults to 2.
+        asam (bool, optional):
+            enables ASAM variant which makes perturbation relative to weight magnitudes.
+            ASAM requires a much larger :code:`rho`, like 0.5 or 1.
+            The :code:`tz.m.ASAM` class is idential to setting this argument to True, but
+            it has larger :code:`rho` by default.
+
+    Examples:
+        SAM-SGD:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.SAM(),
+                tz.m.LR(1e-2)
+            )
+
+        SAM-Adam:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.SAM(),
+                tz.m.Adam(),
+                tz.m.LR(1e-2)
+            )
+
+    References:
+        Foret, P., Kleiner, A., Mobahi, H., & Neyshabur, B. (2020). Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412. https://arxiv.org/abs/2010.01412#page=3.16
+    """
+    def __init__(self, rho: float = 0.05, p: float = 2, eps=1e-10, asam=False):
+        defaults = dict(rho=rho, p=p, eps=eps, asam=asam)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+
+        params = var.params
+        closure = var.closure
+        zero_grad = var.zero_grad
+        if closure is None: raise RuntimeError("SAM requires a closure passed to the optimizer step")
+        p, rho = self.get_settings(var.params, 'p', 'rho', cls=NumberList)
+        s = self.settings[var.params[0]]
+        eps = s['eps']
+        asam = s['asam']
+
+        # 1/p + 1/q = 1
+        # okay, authors of SAM paper, I will manually solve your equation
+        # so q = -p/(1-p)
+        q = -p / (1-p)
+        # as a validation for 2 it is -2 / -1 = 2
+
+        @torch.no_grad
+        def sam_closure(backward=True):
+            orig_grads = None
+            if not backward:
+                # if backward is False, make sure this doesn't modify gradients
+                # to avoid issues
+                orig_grads = [p.grad for p in params]
+
+            # gradient at initial parameters
+            zero_grad()
+            with torch.enable_grad():
+                closure()
+
+            grad = TensorList(p.grad if p.grad is not None else torch.zeros_like(p) for p in params)
+            grad_abs = grad.abs()
+
+            # compute e
+            term1 = grad.sign().mul_(rho)
+            term2 = grad_abs.pow(q-1)
+
+            if asam:
+                grad_abs.mul_(torch._foreach_abs(params))
+
+            denom = grad_abs.pow_(q).sum().pow(1/p)
+
+            e = term1.mul_(term2).div_(denom.clip(min=eps))
+
+            if asam:
+                e.mul_(torch._foreach_pow(params, 2))
+
+            # calculate loss and gradient approximation of inner problem
+            torch._foreach_add_(params, e)
+            if backward:
+                zero_grad()
+                with torch.enable_grad():
+                    # this sets .grad attributes
+                    sam_loss = closure()
+
+            else:
+                sam_loss = closure(False)
+
+            # and restore initial parameters
+            torch._foreach_sub_(params, e)
+
+            if orig_grads is not None:
+                for param,orig_grad in zip(params, orig_grads):
+                    param.grad = orig_grad
+
+            return sam_loss
+
+        var.closure = sam_closure
+        return var
+
+# different class because defaults for SAM are bad for ASAM
+class ASAM(SAM):
+    """Adaptive Sharpness-Aware Minimization from https://arxiv.org/pdf/2102.11600#page=6.52
+
+    SAM functions by seeking parameters that lie in neighborhoods having uniformly low loss value.
+    It performs two forward and backward passes per step.
+
+    This implementation modifies the closure to return loss and calculate gradients
+    of the SAM objective. All modules after this will use the modified objective.
+
+    .. note::
+        This module requires a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients at two points on each step.
+
+    Args:
+        rho (float, optional): Neighborhood size. Defaults to 0.05.
+        p (float, optional): norm of the SAM objective. Defaults to 2.
+
+    Examples:
+        ASAM-Adam:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ASAM(),
+                tz.m.Adam(),
+                tz.m.LR(1e-2)
+            )
+
+    References:
+        Kwon, J., Kim, J., Park, H., & Choi, I. K. (2021, July). Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning (pp. 5905-5914). PMLR. https://arxiv.org/abs/2102.11600
+    """
+    def __init__(self, rho: float = 0.5, p: float = 2, eps=1e-10):
+        super().__init__(rho=rho, p=p, eps=eps, asam=True)

torchzero/modules/optimizers/shampoo.py

@@ -59,7 +59,7 @@ def _merge_small_dims(tensor: torch.Tensor, max_dim: int):
     if tensor.shape[sort_idxs[0]] > max_dim:
         return tensor, None, None
 
-    tensor = tensor.permute(*sort_idxs)
+    tensor = tensor.permute(*sort_idxs.tolist())
     flatten_end_idx = 0
     flat_sizes = []
     flat_numel = 1
@@ -80,19 +80,28 @@ def _unmerge_small_dims(tensor: torch.Tensor, flat_sizes: Sequence[int] | None,
     if flat_sizes is None: return tensor
     assert sort_idxs is not None
     tensor = tensor.unflatten(0, flat_sizes)
-    return tensor.permute(*np.argsort(sort_idxs))
+    return tensor.permute(*np.argsort(sort_idxs).tolist())
 
 
 class Shampoo(Transform):
     """Shampoo from Preconditioned Stochastic Tensor Optimization (https://arxiv.org/abs/1802.09568).
 
+    .. note::
+        Shampoo is usually grafted to another optimizer like Adam, otherwise it can be unstable. An example of how to do grafting is given below in the Examples section.
+
+    .. note::
+        Shampoo is a very computationally expensive optimizer, increase :code:`update_freq` if it is too slow.
+
+    .. note::
+        SOAP optimizer usually outperforms Shampoo and is also not as computationally expensive. SOAP implementation is available as :code:`tz.m.SOAP`.
+
     Args:
         decay (float | None, optional): slowly decays preconditioners. Defaults to None.
         beta (float | None, optional):
             if None calculates sum as in standard shampoo, otherwise uses EMA of preconditioners. Defaults to None.
         matrix_eps (float, optional): epsilon for matrix operations. Defaults to 1e-10.
         update_freq (int, optional): preconditioner update frequency. Defaults to 10.
-        exp_override (int | None, optional): matrix exponent override, if not set, uses 2*ndim. Defaults to None.
+        exp_override (int | None, optional): matrix exponent override, if not set, uses 2*ndim. Defaults to 2.
         merge_small (bool, optional): whether to merge small dims on tensors. Defaults to True.
         max_dim (int, optional): maximum dimension size for preconditioning. Defaults to 2_000.
         precondition_1d (bool, optional): whether to precondition 1d tensors. Defaults to True.
@@ -101,13 +110,38 @@ class Shampoo(Transform):
             module applied after updating preconditioners and before applying preconditioning.
             For example if beta≈0.999 and `inner=tz.m.EMA(0.9)`, this becomes Adam with shampoo preconditioner (ignoring debiasing).
             Defaults to None.
+
+    Examples:
+        Shampoo grafted to Adam
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.GraftModules(
+                    direction = tz.m.Shampoo(),
+                    magnitude = tz.m.Adam(),
+                ),
+                tz.m.LR(1e-3)
+            )
+
+        Adam with Shampoo preconditioner
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Shampoo(beta=0.999, inner=tz.m.EMA(0.9)),
+                tz.m.Debias(0.9, 0.999),
+                tz.m.LR(1e-3)
+            )
     """
     def __init__(
         self,
         decay: float | None = None,
         beta: float | None = None,
         update_freq: int = 10,
-        exp_override: int | None = None,
+        exp_override: int | None = 2,
         merge_small: bool = True,
         max_dim: int = 2_000,
         precondition_1d: bool = True,
@@ -120,7 +154,7 @@ class Shampoo(Transform):
         if inner is not None:
             self.set_child('inner', inner)
 
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         merged_tensors = [] # target with merged dims
 
         # update preconditioners