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

torchzero/modules/misc/misc.py

@@ -0,0 +1,316 @@
+from collections import deque
+from collections.abc import Iterable
+from operator import itemgetter
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, Target, TensorwiseTransform, Transform, Var
+from ...utils import Distributions, NumberList, TensorList, unpack_dicts, unpack_states
+
+
+class Previous(TensorwiseTransform):
+    """Maintains an update from n steps back, for example if n=1, returns previous update"""
+    def __init__(self, n=1, target: Target = 'update'):
+        defaults = dict(n=n)
+        super().__init__(uses_grad=False, defaults=defaults, target=target)
+
+
+    @torch.no_grad
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        n = setting['n']
+
+        if 'history' not in state:
+            state['history'] = deque(maxlen=n+1)
+
+        state['history'].append(tensor)
+
+        return state['history'][0]
+
+
+class LastDifference(Transform):
+    """Outputs difference between past two updates."""
+    def __init__(self,target: Target = 'update'):
+        super().__init__({}, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        prev_tensors = unpack_states(states, tensors, 'prev_tensors') # initialized to 0
+        difference = torch._foreach_sub(tensors, prev_tensors)
+        for p, c in zip(prev_tensors, tensors): p.set_(c)
+        return difference
+
+class LastGradDifference(Module):
+    """Outputs difference between past two gradients."""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def step(self, var):
+        grad = var.get_grad()
+        prev_grad = self.get_state(var.params, 'prev_grad') # initialized to 0
+        difference = torch._foreach_sub(grad, prev_grad)
+        for p, c in zip(prev_grad, grad): p.copy_(c)
+        var.update = list(difference)
+        return var
+
+class LastParamDifference(Module):
+    """Outputs difference between past two parameters, which is the effective previous update."""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def step(self, var):
+        params = var.params
+        prev_params = self.get_state(var.params, 'prev_params') # initialized to 0
+        difference = torch._foreach_sub(params, prev_params)
+        for p, c in zip(prev_params, params): p.copy_(c)
+        var.update = list(difference)
+        return var
+
+
+
+class LastProduct(Transform):
+    """Outputs difference between past two updates."""
+    def __init__(self,target: Target = 'update'):
+        super().__init__({}, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', init=torch.ones_like) # initialized to 1 for prod
+        prod = torch._foreach_mul(tensors, prev)
+        for p, c in zip(prev, tensors): p.set_(c)
+        return prod
+
+class LastRatio(Transform):
+    """Outputs ratio between past two updates, the numerator is determined by :code:`numerator` argument."""
+    def __init__(self, numerator: Literal['cur', 'prev'] = 'cur', target: Target = 'update'):
+        defaults = dict(numerator=numerator)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', init = torch.ones_like) # initialized to ones
+        numerator = settings[0]['numerator']
+        if numerator == 'cur': ratio = torch._foreach_div(tensors, prev)
+        else: ratio = torch._foreach_div(prev, tensors)
+        for p, c in zip(prev, tensors): p.set_(c)
+        return ratio
+
+class LastAbsoluteRatio(Transform):
+    """Outputs ratio between absolute values of past two updates the numerator is determined by :code:`numerator` argument."""
+    def __init__(self, numerator: Literal['cur', 'prev'] = 'cur', eps:float=1e-8, target: Target = 'update'):
+        defaults = dict(numerator=numerator, eps=eps)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        prev = unpack_states(states, tensors, 'prev', init = torch.ones_like) # initialized to ones
+        numerator = settings[0]['numerator']
+        eps = NumberList(s['eps'] for s in settings)
+
+        torch._foreach_abs_(tensors)
+        torch._foreach_clamp_min_(prev, eps)
+
+        if numerator == 'cur': ratio = torch._foreach_div(tensors, prev)
+        else: ratio = torch._foreach_div(prev, tensors)
+        for p, c in zip(prev, tensors): p.set_(c)
+        return ratio
+
+class GradSign(Transform):
+    """Copies gradient sign to update."""
+    def __init__(self, target: Target = 'update'):
+        super().__init__({}, uses_grad=True, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        assert grads is not None
+        return [t.copysign_(g) for t,g in zip(tensors, grads)]
+
+class UpdateSign(Transform):
+    """Outputs gradient with sign copied from the update."""
+    def __init__(self, target: Target = 'update'):
+        super().__init__({}, uses_grad=True, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        assert grads is not None
+        return [g.copysign(t) for t,g in zip(tensors, grads)] # no in-place
+
+class GraftToGrad(Transform):
+    """Grafts update to the gradient, that is update is rescaled to have the same norm as the gradient."""
+    def __init__(self, tensorwise:bool=False, ord:float=2, eps:float = 1e-6, target: Target = 'update'):
+        defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps)
+        super().__init__(defaults, uses_grad=True, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        assert grads is not None
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(settings[0])
+        return TensorList(tensors).graft_(grads, tensorwise=tensorwise, ord=ord, eps=eps)
+
+class GraftGradToUpdate(Transform):
+    """Outputs gradient grafted to update, that is gradient rescaled to have the same norm as the update."""
+    def __init__(self, tensorwise:bool=False, ord:float=2, eps:float = 1e-6, target: Target = 'update'):
+        defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps)
+        super().__init__(defaults, uses_grad=True, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        assert grads is not None
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(settings[0])
+        return TensorList(grads).graft(tensors, tensorwise=tensorwise, ord=ord, eps=eps)
+
+
+class GraftToParams(Transform):
+    """Grafts update to the parameters, that is update is rescaled to have the same norm as the parameters, but no smaller than :code:`eps`."""
+    def __init__(self, tensorwise:bool=False, ord:float=2, eps:float = 1e-4, target: Target = 'update'):
+        defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensorwise, ord, eps = itemgetter('tensorwise','ord','eps')(settings[0])
+        return TensorList(tensors).graft_(params, tensorwise=tensorwise, ord=ord, eps=eps)
+
+class Relative(Transform):
+    """Multiplies update by absolute parameter values to make it relative to their magnitude, :code:`min_value` is minimum allowed value to avoid getting stuck at 0."""
+    def __init__(self, min_value:float = 1e-4, target: Target = 'update'):
+        defaults = dict(min_value=min_value)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        mul = TensorList(params).abs().clamp_([s['min_value'] for s in settings])
+        torch._foreach_mul_(tensors, mul)
+        return tensors
+
+class FillLoss(Module):
+    """Outputs tensors filled with loss value times :code:`alpha`"""
+    def __init__(self, alpha: float = 1, backward: bool = True):
+        defaults = dict(alpha=alpha, backward=backward)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        alpha = self.get_settings(var.params, 'alpha')
+        loss = var.get_loss(backward=self.settings[var.params[0]]['backward'])
+        var.update = [torch.full_like(p, loss*a) for p,a in zip(var.params, alpha)]
+        return var
+
+class MulByLoss(Module):
+    """Multiplies update by loss times :code:`alpha`"""
+    def __init__(self, alpha: float = 1, min_value:float = 1e-8, backward: bool = True):
+        defaults = dict(alpha=alpha, min_value=min_value, backward=backward)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        alpha, min_value = self.get_settings(var.params, 'alpha', 'min_value')
+        loss = var.get_loss(backward=self.settings[var.params[0]]['backward'])
+        mul = [max(loss*a, mv) for a,mv in zip(alpha, min_value)]
+        torch._foreach_mul_(var.update, mul)
+        return var
+
+class DivByLoss(Module):
+    """Divides update by loss times :code:`alpha`"""
+    def __init__(self, alpha: float = 1, min_value:float = 1e-8, backward: bool = True):
+        defaults = dict(alpha=alpha, min_value=min_value, backward=backward)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        alpha, min_value = self.get_settings(var.params, 'alpha', 'min_value')
+        loss = var.get_loss(backward=self.settings[var.params[0]]['backward'])
+        mul = [max(loss*a, mv) for a,mv in zip(alpha, min_value)]
+        torch._foreach_div_(var.update, mul)
+        return var
+
+
+class NoiseSign(Transform):
+    """Outputs random tensors with sign copied from the update."""
+    def __init__(self, distribution:Distributions = 'normal', alpha = 1):
+        defaults = dict(distribution=distribution, alpha=alpha)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        alpha = [s['alpha'] for s in settings]
+        distribution = self.settings[params[0]]['distribution']
+        return TensorList(tensors).sample_like(alpha, distribution).copysign_(tensors)
+
+class HpuEstimate(Transform):
+    """returns ``y/||s||``, where ``y`` is difference between current and previous update (gradient), ``s`` is difference between current and previous parameters. The returned tensors are a finite difference approximation to hessian times previous update."""
+    def __init__(self):
+        defaults = dict()
+        super().__init__(defaults, uses_grad=False)
+
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_params', 'prev_update')
+
+    @torch.no_grad
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
+        prev_params, prev_update = self.get_state(params, 'prev_params', 'prev_update') # initialized to 0
+        s = torch._foreach_sub(params, prev_params)
+        y = torch._foreach_sub(tensors, prev_update)
+        for p, c in zip(prev_params, params): p.copy_(c)
+        for p, c in zip(prev_update, tensors): p.copy_(c)
+        torch._foreach_div_(y, torch.linalg.norm(torch.cat([t.ravel() for t in s])).clip(min=1e-8)) # pylint:disable=not-callable
+        self.store(params, ['s', 'y'], [s, y])
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        return [self.state[p]['y'] for p in params]
+
+class RandomHvp(Module):
+    """Returns a hessian-vector product with a random vector"""
+
+    def __init__(
+        self,
+        n_samples: int = 1,
+        distribution: Distributions = "normal",
+        update_freq: int = 1,
+        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
+        h=1e-3,
+    ):
+        defaults = dict(n_samples=n_samples, distribution=distribution, hvp_method=hvp_method, h=h, update_freq=update_freq)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        params = TensorList(var.params)
+        settings = self.settings[params[0]]
+        n_samples = settings['n_samples']
+        distribution = settings['distribution']
+        hvp_method = settings['hvp_method']
+        h = settings['h']
+        update_freq = settings['update_freq']
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        D = None
+        if step % update_freq == 0:
+
+            rgrad = None
+            for i in range(n_samples):
+                u = params.sample_like(distribution=distribution)
+
+                Hvp, rgrad = self.Hvp(u, at_x0=True, var=var, rgrad=rgrad, hvp_method=hvp_method,
+                                    h=h, normalize=True, retain_grad=i < n_samples-1)
+
+                if D is None: D = Hvp
+                else: torch._foreach_add_(D, Hvp)
+
+            if n_samples > 1: torch._foreach_div_(D, n_samples)
+            if update_freq != 1:
+                assert D is not None
+                D_buf = self.get_state(params, "D", cls=TensorList)
+                D_buf.set_(D)
+
+        if D is None:
+            D = self.get_state(params, "D", cls=TensorList)
+
+        var.update = list(D)
+        return var

torchzero/modules/misc/multistep.py

@@ -0,0 +1,158 @@
+from collections.abc import Iterable
+
+import torch
+
+from ...core import Chainable, Module, Var
+from ...utils import TensorList
+
+def _sequential_step(self: Module, var: Var, sequential: bool):
+    params = var.params
+    steps = self.settings[params[0]]['steps']
+
+    if sequential: modules = self.get_children_sequence() * steps
+    else: modules = [self.children['module']] * steps
+
+    if var.closure is None and len(modules) > 1: raise ValueError('Multistep and Sequential require closure')
+
+    # store original params unless this is last module and can update params directly
+    params_before_steps = None if (var.is_last and var.last_module_lrs is None) else [p.clone() for p in params]
+
+    # first step - pass var as usual
+    var = modules[0].step(var)
+    new_var = var
+
+    # subsequent steps - update parameters and create new var
+    if len(modules) > 1:
+        for m in modules[1:]:
+
+            # update params
+            if (not new_var.skip_update):
+                if new_var.last_module_lrs is not None:
+                    torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
+
+                torch._foreach_sub_(params, new_var.get_update())
+
+            # create new var since we are at a new point, that means grad, update and loss will be None
+            new_var = Var(params=new_var.params, closure=new_var.closure,
+                            model=new_var.model, current_step=new_var.current_step + 1)
+
+            # step
+            new_var = m.step(new_var)
+
+        # final parameter update
+        if (not new_var.skip_update):
+            if new_var.last_module_lrs is not None:
+                torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
+
+            torch._foreach_sub_(params, new_var.get_update())
+
+    # if last module, update is applied so return new var
+    if params_before_steps is None:
+        new_var.stop = True
+        new_var.skip_update = True
+        return new_var
+
+    # otherwise use parameter difference as update
+    var.update = list(torch._foreach_sub(params_before_steps, params))
+    for p, bef in zip(params, params_before_steps):
+        p.set_(bef) # pyright:ignore[reportArgumentType]
+    return var
+
+class Multistep(Module):
+    """Performs :code:`steps` inner steps with :code:`module` per each step.
+
+    The update is taken to be the parameter difference between parameters before and after the inner loop."""
+    def __init__(self, module: Chainable, steps: int):
+        defaults = dict(steps=steps)
+        super().__init__(defaults)
+        self.set_child('module', module)
+
+    @torch.no_grad
+    def step(self, var):
+        return _sequential_step(self, var, sequential=False)
+
+class Sequential(Module):
+    """On each step, this sequentially steps with :code:`modules` :code:`steps` times.
+
+    The update is taken to be the parameter difference between parameters before and after the inner loop."""
+    def __init__(self, modules: Iterable[Chainable], steps: int=1):
+        defaults = dict(steps=steps)
+        super().__init__(defaults)
+        self.set_children_sequence(modules)
+
+    @torch.no_grad
+    def step(self, var):
+        return _sequential_step(self, var, sequential=True)
+
+
+class NegateOnLossIncrease(Module):
+    """Uses an extra forward pass to evaluate loss at :code:`parameters+update`,
+    if loss is larger than at :code:`parameters`,
+    the update is set to 0 if :code:`backtrack=False` and to :code:`-update` otherwise"""
+    def __init__(self, backtrack=False):
+        defaults = dict(backtrack=backtrack)
+        super().__init__(defaults=defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+        if closure is None: raise RuntimeError('NegateOnLossIncrease requires closure')
+        backtrack = self.settings[var.params[0]]['backtrack']
+
+        update = var.get_update()
+        f_0 = var.get_loss(backward=False)
+
+        torch._foreach_sub_(var.params, update)
+        f_1 = closure(False)
+
+        if f_1 <= f_0:
+            if var.is_last and var.last_module_lrs is None:
+                var.stop = True
+                var.skip_update = True
+                return var
+
+            torch._foreach_add_(var.params, update)
+            return var
+
+        torch._foreach_add_(var.params, update)
+        if backtrack:
+            torch._foreach_neg_(var.update)
+        else:
+            torch._foreach_zero_(var.update)
+        return var
+
+
+class Online(Module):
+    """Allows certain modules to be used for mini-batch optimization."""
+    def __init__(self, module: Chainable,):
+        super().__init__()
+
+        self.set_child('module', module)
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+        if closure is None: raise ValueError("Closure must be passed for Online")
+        step = self.global_state.get('step', 0) + 1
+        self.global_state['step'] = step
+        params = TensorList(var.params)
+        p_cur = params.clone()
+        p_prev = self.get_state(params, 'p_prev', cls=TensorList)
+        module = self.children['module']
+
+        if step == 1:
+            var = module.step(var.clone(clone_update=False))
+
+            p_prev.copy_(params)
+            return var
+
+        # restore previous params
+        var_prev = Var(params=params, closure=closure, model=var.model, current_step=var.current_step)
+        params.set_(p_prev)
+        module.reset_for_online()
+        module.update(var_prev)
+
+        # restore current params
+        params.set_(p_cur)
+        p_prev.copy_(params)
+        return module.step(var.clone(clone_update=False))

torchzero/modules/misc/regularization.py

@@ -0,0 +1,171 @@
+from collections import deque
+from collections.abc import Iterable
+from operator import itemgetter
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, Target, TensorwiseTransform, Transform, Var
+from ...utils import Distributions, NumberList, TensorList, unpack_dicts, unpack_states
+
+
+class Dropout(Transform):
+    """Applies dropout to the update.
+
+    For each weight the update to that weight has :code:`p` probability to be set to 0.
+    This can be used to implement gradient dropout or update dropout depending on placement.
+
+    Args:
+        p (float, optional): probability that update for a weight is replaced with 0. Defaults to 0.5.
+        graft (bool, optional):
+            if True, update after dropout is rescaled to have the same norm as before dropout. Defaults to False.
+        target (Target, optional): what to set on var, refer to documentation. Defaults to 'update'.
+
+
+    Examples:
+        Gradient dropout.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Dropout(0.5),
+                tz.m.Adam(),
+                tz.m.LR(1e-3)
+            )
+
+        Update dropout.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.Dropout(0.5),
+                tz.m.LR(1e-3)
+            )
+
+    """
+    def __init__(self, p: float = 0.5, graft: bool=False, target: Target = 'update'):
+        defaults = dict(p=p, graft=graft)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = TensorList(tensors)
+        p = NumberList(s['p'] for s in settings)
+        graft = settings[0]['graft']
+
+        if graft:
+            target_norm = tensors.global_vector_norm()
+            tensors.mul_(tensors.rademacher_like(1-p).add_(1).div_(2))
+            return tensors.mul_(target_norm / tensors.global_vector_norm()) # graft
+
+        return tensors.mul_(tensors.rademacher_like(1-p).add_(1).div_(2))
+
+def _bernoulli_like(tensor, p = 0.5, generator = None):
+    """p is probability of a 1, other values will be 0."""
+    return torch.bernoulli(torch.full_like(tensor, p), generator = generator)
+
+class WeightDropout(Module):
+    """
+    Changes the closure so that it evaluates loss and gradients with random weights replaced with 0.
+
+    Dropout can be disabled for a parameter by setting :code:`use_dropout=False` in corresponding parameter group.
+
+    Args:
+        p (float, optional): probability that any weight is replaced with 0. Defaults to 0.5.
+        graft (bool, optional):
+            if True, parameters after dropout are rescaled to have the same norm as before dropout. Defaults to False.
+    """
+    def __init__(self, p: float = 0.5, graft: bool = True):
+        defaults = dict(p=p, graft=graft, use_dropout=True)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+        if closure is None: raise RuntimeError('WeightDropout requires closure')
+        params = TensorList(var.params)
+        p = NumberList(self.settings[p]['p'] for p in params)
+
+        # create masks
+        mask = []
+        for p, m in zip(params, mask):
+            prob = self.settings[p]['p']
+            use_dropout = self.settings[p]['use_dropout']
+            if use_dropout: mask.append(_bernoulli_like(p, prob))
+            else: mask.append(torch.ones_like(p))
+
+        @torch.no_grad
+        def dropout_closure(backward=True):
+            orig_params = params.clone()
+            params.mul_(mask)
+            if backward:
+                with torch.enable_grad(): loss = closure()
+            else:
+                loss = closure(False)
+            params.copy_(orig_params)
+            return loss
+
+        var.closure = dropout_closure
+        return var
+
+
+class PerturbWeights(Module):
+    """
+    Changes the closure so that it evaluates loss and gradients at weights perturbed by a random perturbation.
+
+    Can be disabled for a parameter by setting :code:`perturb=False` in corresponding parameter group.
+
+    Args:
+        alpha (float, optional): multiplier for perturbation magnitude. Defaults to 0.1.
+        relative (bool, optional): whether to multiply perturbation by mean absolute value of the parameter. Defaults to True.
+        graft (bool, optional):
+            if True, parameters after dropout are rescaled to have the same norm as before dropout. Defaults to False.
+    """
+    def __init__(self, alpha: float = 0.1, relative:bool=True, distribution:Distributions = 'normal'):
+        defaults = dict(alpha=alpha, relative=relative, distribution=distribution, perturb=True)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+        if closure is None: raise RuntimeError('WeightDropout requires closure')
+        params = TensorList(var.params)
+
+        # create perturbations
+        perts = []
+        for p in params:
+            settings = self.settings[p]
+            if not settings['perturb']:
+                perts.append(torch.zeros_like(p))
+                continue
+
+            alpha = settings['alpha']
+            if settings['relative']:
+                alpha *= p.abs().mean()
+
+            distribution = self.settings[p]['distribution'].lower()
+            if distribution in ('normal', 'gaussian'):
+                perts.append(torch.randn_like(p).mul_(alpha))
+            elif distribution == 'uniform':
+                perts.append(torch.empty_like(p).uniform_(-alpha,alpha))
+            elif distribution == 'sphere':
+                r = torch.randn_like(p)
+                perts.append((r * alpha) / torch.linalg.vector_norm(r)) # pylint:disable=not-callable
+            else:
+                raise ValueError(distribution)
+
+        @torch.no_grad
+        def perturbed_closure(backward=True):
+            params.add_(perts)
+            if backward:
+                with torch.enable_grad(): loss = closure()
+            else:
+                loss = closure(False)
+            params.sub_(perts)
+            return loss
+
+        var.closure = perturbed_closure
+        return var

torchzero/modules/{ops → misc}/split.py

@@ -45,7 +45,35 @@ def _split(
     return var
 
 class Split(Module):
-    """Apply `true` modules to all parameters filtered by `filter`, apply `false` modules to all other parameters."""
+    """Apply `true` modules to all parameters filtered by `filter`, apply `false` modules to all other parameters.
+
+    Args:
+        filter (Callable[[torch.Tensor], bool]): a function that takes in a parameter tensor and returns a boolean value.
+        true (Chainable | None): modules that are applied to tensors where :code:`filter` returned True.
+        false (Chainable | None): modules that are applied to tensors where :code:`filter` returned False.
+
+    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)
+            )
+
+
+    """
     def __init__(self, filter: Callable[[torch.Tensor], bool], true: Chainable | None, false: Chainable | None):
         defaults = dict(filter=filter)
         super().__init__(defaults)