torchzero 0.1.7__py3-none-any.whl → 0.3.1__py3-none-any.whl

torchzero/modules/ops/reduce.py

@@ -0,0 +1,149 @@
+""""""
+from abc import ABC, abstractmethod
+from collections.abc import Iterable,Sequence
+from typing import Any, cast
+
+import torch
+
+from ...core import Chainable, Module, Target, Vars, maybe_chain
+
+
+class ReduceOperation(Module, ABC):
+    """Base class for reduction operations like Sum, Prod, Maximum. This is an abstract class, subclass it and override `transform` method to use it."""
+    def __init__(self, defaults: dict[str, Any] | None, *operands: Chainable | Any):
+        super().__init__(defaults=defaults)
+
+        self.operands = []
+        for i, v in enumerate(operands):
+
+            if isinstance(v, (Module, Sequence)):
+                self.set_child(f'operand_{i}', v)
+                self.operands.append(self.children[f'operand_{i}'])
+            else:
+                self.operands.append(v)
+
+        if not self.children:
+            raise ValueError('At least one operand must be a module')
+
+    @abstractmethod
+    def transform(self, vars: Vars, *operands: Any | list[torch.Tensor]) -> list[torch.Tensor]:
+        """applies the operation to operands"""
+        raise NotImplementedError
+
+    @torch.no_grad
+    def step(self, vars: Vars) -> Vars:
+        # pass cloned update to all module operands
+        processed_operands: list[Any | list[torch.Tensor]] = self.operands.copy()
+
+        for i, v in enumerate(self.operands):
+            if f'operand_{i}' in self.children:
+                v: Module
+                updated_vars = v.step(vars.clone(clone_update=True))
+                processed_operands[i] = updated_vars.get_update()
+                vars.update_attrs_from_clone_(updated_vars) # update loss, grad, etc if this module calculated them
+
+        transformed = self.transform(vars, *processed_operands)
+        vars.update = transformed
+        return vars
+
+class Sum(ReduceOperation):
+    USE_MEAN = False
+    def __init__(self, *inputs: Chainable | float):
+        super().__init__({}, *inputs)
+
+    @torch.no_grad
+    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+        sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
+        sum = cast(list, sorted_inputs[0])
+        if len(sorted_inputs) > 1:
+            for v in sorted_inputs[1:]:
+                torch._foreach_add_(sum, v)
+
+        if self.USE_MEAN and len(sorted_inputs) > 1: torch._foreach_div_(sum, len(sorted_inputs))
+        return sum
+
+class Mean(Sum):
+    USE_MEAN = True
+
+
+class WeightedSum(ReduceOperation):
+    USE_MEAN = False
+    def __init__(self, *inputs: Chainable | float, weights: Iterable[float]):
+        weights = list(weights)
+        if len(inputs) != len(weights):
+            raise ValueError(f'Number of inputs {len(inputs)} must match number of weights {len(weights)}')
+        defaults = dict(weights=weights)
+        super().__init__(defaults=defaults, *inputs)
+
+    @torch.no_grad
+    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+        sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
+        weights = self.settings[vars.params[0]]['weights']
+        sum = cast(list, sorted_inputs[0])
+        torch._foreach_mul_(sum, weights[0])
+        if len(sorted_inputs) > 1:
+            for v, w in zip(sorted_inputs[1:], weights[1:]):
+                if isinstance(v, (int, float)): torch._foreach_add_(sum, v*w)
+                else: torch._foreach_add_(sum, v, alpha=w)
+
+        if self.USE_MEAN and len(sorted_inputs) > 1: torch._foreach_div_(sum, len(sorted_inputs))
+        return sum
+
+
+class WeightedMean(WeightedSum):
+    USE_MEAN = True
+
+class Median(ReduceOperation):
+    def __init__(self, *inputs: Chainable | float):
+        super().__init__({}, *inputs)
+
+    @torch.no_grad
+    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+        res = []
+        lists = [i for i in inputs if isinstance(i, list)]
+        floats = [i for i in inputs if isinstance(i, (int,float))]
+        for tensors in zip(*lists):
+            res.append(torch.median(torch.stack(tensors + tuple(torch.full_like(tensors[0], f) for f in floats)), dim=0))
+        return res
+
+class Prod(ReduceOperation):
+    def __init__(self, *inputs: Chainable | float):
+        super().__init__({}, *inputs)
+
+    @torch.no_grad
+    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+        sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
+        prod = cast(list, sorted_inputs[0])
+        if len(sorted_inputs) > 1:
+            for v in sorted_inputs[1:]:
+                torch._foreach_mul_(prod, v)
+
+        return prod
+
+class MaximumModules(ReduceOperation):
+    def __init__(self, *inputs: Chainable | float):
+        super().__init__({}, *inputs)
+
+    @torch.no_grad
+    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+        sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
+        maximum = cast(list, sorted_inputs[0])
+        if len(sorted_inputs) > 1:
+            for v in sorted_inputs[1:]:
+                torch._foreach_maximum_(maximum, v)
+
+        return maximum
+
+class MinimumModules(ReduceOperation):
+    def __init__(self, *inputs: Chainable | float):
+        super().__init__({}, *inputs)
+
+    @torch.no_grad
+    def transform(self, vars: Vars, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
+        sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
+        minimum = cast(list, sorted_inputs[0])
+        if len(sorted_inputs) > 1:
+            for v in sorted_inputs[1:]:
+                torch._foreach_minimum_(minimum, v)
+
+        return minimum

torchzero/modules/ops/split.py

@@ -0,0 +1,75 @@
+from collections.abc import Callable
+from typing import cast
+
+import torch
+
+from ...core import Chainable, Module, Vars
+
+
+def _split(
+    module: Module,
+    idxs,
+    params,
+    vars: Vars,
+):
+    split_params = [p for i,p in enumerate(params) if i in idxs]
+
+    split_grad = None
+    if vars.grad is not None:
+        split_grad = [g for i,g in enumerate(vars.grad) if i in idxs]
+
+    split_update = None
+    if vars.update is not None:
+        split_update = [u for i,u in enumerate(vars.update) if i in idxs]
+
+    split_vars = vars.clone(clone_update=False)
+    split_vars.params = split_params
+    split_vars.grad = split_grad
+    split_vars.update = split_update
+
+    split_vars = module.step(split_vars)
+
+    if (vars.grad is None) and (split_vars.grad is not None):
+        vars.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+
+    if split_vars.update is not None:
+
+        if vars.update is None:
+            if vars.grad is None: vars.update = [cast(torch.Tensor, None) for _ in vars.params]
+            else: vars.update = [g.clone() for g in vars.grad]
+
+        for idx, u in zip(idxs, split_vars.update):
+            vars.update[idx] = u
+
+    vars.update_attrs_from_clone_(split_vars)
+    return vars
+
+class Split(Module):
+    """Apply `true` modules to all parameters filtered by `filter`, apply `false` modules to all other parameters."""
+    def __init__(self, filter: Callable[[torch.Tensor], bool], true: Chainable | None, false: Chainable | None):
+        defaults = dict(filter=filter)
+        super().__init__(defaults)
+
+        if true is not None: self.set_child('true', true)
+        if false is not None: self.set_child('false', false)
+
+    def step(self, vars):
+
+        params = vars.params
+        filter = self.settings[params[0]]['filter']
+
+        true_idxs = []
+        false_idxs = []
+        for i,p in enumerate(params):
+            if filter(p): true_idxs.append(i)
+            else: false_idxs.append(i)
+
+        if 'true' in self.children:
+            true = self.children['true']
+            vars = _split(true, idxs=true_idxs, params=params, vars=vars)
+
+        if 'false' in self.children:
+            false = self.children['false']
+            vars = _split(false, idxs=false_idxs, params=params, vars=vars)
+
+        return vars

torchzero/modules/ops/switch.py

@@ -0,0 +1,68 @@
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+import torch
+
+from ...core import Chainable, Module
+
+
+class Alternate(Module):
+    """alternate between stepping with `modules`"""
+    LOOP = True
+    def __init__(self, *modules: Chainable, steps: int | Iterable[int] = 1):
+        if isinstance(steps, Iterable):
+            steps = list(steps)
+            if len(steps) != len(modules):
+                raise ValueError(f"steps must be the same length as modules, got {len(modules) = }, {len(steps) = }")
+
+        defaults = dict(steps=steps)
+        super().__init__(defaults)
+
+        self.set_children_sequence(modules)
+        self.global_state['current_module_idx'] = 0
+        self.global_state['steps_to_next'] = steps[0] if isinstance(steps, list) else steps
+
+    @torch.no_grad
+    def step(self, vars):
+        # get current module
+        current_module_idx = self.global_state.setdefault('current_module_idx', 0)
+        module = self.children[f'module_{current_module_idx}']
+
+        # step
+        vars = module.step(vars.clone(clone_update=False))
+
+        # number of steps until next module
+        steps = self.settings[vars.params[0]]['steps']
+        if isinstance(steps, int): steps = [steps]*len(self.children)
+
+        if 'steps_to_next' not in self.global_state:
+            self.global_state['steps_to_next'] = steps[0] if isinstance(steps, list) else steps
+
+        self.global_state['steps_to_next'] -= 1
+
+        # switch to next module
+        if self.global_state['steps_to_next'] == 0:
+            self.global_state['current_module_idx'] += 1
+
+            # loop to first module (or keep using last module on Switch)
+            if self.global_state['current_module_idx'] > len(self.children) - 1:
+                if self.LOOP: self.global_state['current_module_idx'] = 0
+                else: self.global_state['current_module_idx'] = len(self.children) - 1
+
+            self.global_state['steps_to_next'] = steps[self.global_state['current_module_idx']]
+
+        return vars
+
+class Switch(Alternate):
+    """switch to next module after some steps"""
+    LOOP = False
+    def __init__(self, *modules: Chainable, steps: int | Iterable[int]):
+
+        if isinstance(steps, Iterable):
+            steps = list(steps)
+            if len(steps) != len(modules) - 1:
+                raise ValueError(f"steps must be the same length as modules, got {len(modules) = }, {len(steps) = }")
+
+            steps.append(1)
+
+        super().__init__(*modules, steps=steps)

torchzero/modules/ops/unary.py

@@ -0,0 +1,115 @@
+from collections import deque
+
+import torch
+
+from ...core import TensorwiseTransform, Target, Transform
+from ...utils import TensorList
+
+class UnaryLambda(Transform):
+    def __init__(self, fn, target: "Target" = 'update'):
+        defaults = dict(fn=fn)
+        super().__init__(defaults=defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        return self.settings[params[0]]['fn'](tensors)
+
+class UnaryParameterwiseLambda(TensorwiseTransform):
+    def __init__(self, fn, target: "Target" = 'update'):
+        defaults = dict(fn=fn)
+        super().__init__(uses_grad=False, defaults=defaults, target=target)
+
+    @torch.no_grad
+    def transform(self, tensor, param, grad, vars):
+        return self.settings[param]['fn'](tensor)
+
+class CustomUnaryOperation(Transform):
+    def __init__(self, name: str, target: "Target" = 'update'):
+        defaults = dict(name=name)
+        super().__init__(defaults=defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        return getattr(tensors, self.settings[params[0]]['name'])()
+
+
+class Abs(Transform):
+    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        torch._foreach_abs_(tensors)
+        return tensors
+
+class Sign(Transform):
+    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        torch._foreach_sign_(tensors)
+        return tensors
+
+class Exp(Transform):
+    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        torch._foreach_exp_(tensors)
+        return tensors
+
+class Sqrt(Transform):
+    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        torch._foreach_sqrt_(tensors)
+        return tensors
+
+class Reciprocal(Transform):
+    def __init__(self, eps = 0, target: "Target" = 'update'):
+        defaults = dict(eps = eps)
+        super().__init__(defaults, uses_grad=False, target=target)
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        eps = self.get_settings('eps', params=params)
+        if any(e != 0 for e in eps): torch._foreach_add_(tensors, eps)
+        torch._foreach_reciprocal_(tensors)
+        return tensors
+
+class Negate(Transform):
+    def __init__(self, target: "Target" = 'update'): super().__init__({}, uses_grad=False, target=target)
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        torch._foreach_neg_(tensors)
+        return tensors
+
+
+class NanToNum(Transform):
+    """Convert `nan`, `inf` and `-inf` to numbers.
+
+    Args:
+        nan (optional): the value to replace NaNs with. Default is zero.
+        posinf (optional): if a Number, the value to replace positive infinity values with.
+            If None, positive infinity values are replaced with the greatest finite value
+            representable by input's dtype. Default is None.
+        neginf (optional): if a Number, the value to replace negative infinity values with.
+            If None, negative infinity values are replaced with the lowest finite value
+            representable by input's dtype. Default is None.
+    """
+    def __init__(self, nan=None, posinf=None, neginf=None, target: "Target" = 'update'):
+        defaults = dict(nan=nan, posinf=posinf, neginf=neginf)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        nan, posinf, neginf = self.get_settings('nan', 'posinf', 'neginf', params=params)
+        return [t.nan_to_num_(nan_i, posinf_i, neginf_i) for t, nan_i, posinf_i, neginf_i in zip(tensors, nan, posinf, neginf)]
+
+class Rescale(Transform):
+    """rescale update to (min, max) range"""
+    def __init__(self, min: float, max: float, tensorwise: bool = False, eps:float=1e-8, target: "Target" = 'update'):
+        defaults = dict(min=min, max=max, eps=eps, tensorwise=tensorwise)
+        super().__init__(defaults, uses_grad=False, target=target)
+
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars):
+        min,max = self.get_settings('min','max', params=params)
+        tensorwise = self.settings[params[0]]['tensorwise']
+        dim = None if tensorwise else 'global'
+        return TensorList(tensors).rescale(min=min, max=max, eps=self.settings[params[0]]['eps'], dim=dim)

torchzero/modules/ops/utility.py

@@ -0,0 +1,112 @@
+from collections import deque
+
+import torch
+
+from ...core import Module, Target, Transform
+from ...utils.tensorlist import Distributions, TensorList
+
+
+class Clone(Transform):
+    def __init__(self): super().__init__({}, uses_grad=False)
+    @torch.no_grad
+    def transform(self, tensors, params, grads, vars): return [t.clone() for t in tensors]
+
+class Grad(Module):
+    def __init__(self):
+        super().__init__({})
+    @torch.no_grad
+    def step(self, vars):
+        vars.update = [g.clone() for g in vars.get_grad()]
+        return vars
+
+class Params(Module):
+    def __init__(self):
+        super().__init__({})
+    @torch.no_grad
+    def step(self, vars):
+        vars.update = [p.clone() for p in vars.params]
+        return vars
+
+class Update(Module):
+    def __init__(self):
+        super().__init__({})
+    @torch.no_grad
+    def step(self, vars):
+        vars.update = [u.clone() for u in vars.get_update()]
+        return vars
+
+class Zeros(Module):
+    def __init__(self):
+        super().__init__({})
+    @torch.no_grad
+    def step(self, vars):
+        vars.update = [torch.zeros_like(p) for p in vars.params]
+        return vars
+
+class Ones(Module):
+    def __init__(self):
+        super().__init__({})
+    @torch.no_grad
+    def step(self, vars):
+        vars.update = [torch.ones_like(p) for p in vars.params]
+        return vars
+
+class Fill(Module):
+    def __init__(self, value: float):
+        defaults = dict(value=value)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, vars):
+        vars.update = [torch.full_like(p, self.settings[p]['value']) for p in vars.params]
+        return vars
+
+class RandomSample(Module):
+    def __init__(self, eps: float = 1, distribution: Distributions = 'normal'):
+        defaults = dict(eps=eps, distribution=distribution)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, vars):
+        vars.update = TensorList(vars.params).sample_like(
+            eps=self.get_settings('eps',params=vars.params), distribution=self.settings[vars.params[0]]['distribution']
+        )
+        return vars
+
+class Randn(Module):
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def step(self, vars):
+        vars.update = [torch.randn_like(p) for p in vars.params]
+        return vars
+
+class Uniform(Module):
+    def __init__(self, low: float, high: float):
+        defaults = dict(low=low, high=high)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, vars):
+        low,high = self.get_settings('low','high', params=vars.params)
+        vars.update = [torch.empty_like(t).uniform_(l,h) for t,l,h in zip(vars.params, low, high)]
+        return vars
+
+class GradToNone(Module):
+    def __init__(self): super().__init__()
+    def step(self, vars):
+        vars.grad = None
+        return vars
+
+class UpdateToNone(Module):
+    def __init__(self): super().__init__()
+    def step(self, vars):
+        vars.update = None
+        return vars
+
+class Identity(Module):
+    def __init__(self, *args, **kwargs): super().__init__()
+    def step(self, vars): return vars
+
+NoOp = Identity

torchzero/modules/optimizers/__init__.py

@@ -1,10 +1,18 @@
-r"""
-This include various optimizers as composable modules.
-"""
-# from .adam import Adam
-from .sgd import SGD
-from .rprop import Rprop
-from .rmsprop import RMSProp
-from .adagrad import Adagrad
-from .adam import Adam
-from .lion import Lion
+from .adagrad import Adagrad, FullMatrixAdagrad
+from .adam import Adam
+from .lion import Lion
+from .muon import DualNormCorrection, MuonAdjustLR, Orthogonalize, orthogonalize_grads_
+from .rmsprop import RMSprop
+from .rprop import (
+    BacktrackOnSignChange,
+    Rprop,
+    ScaleLRBySignChange,
+    SignConsistencyLRs,
+    SignConsistencyMask,
+)
+from .shampoo import Shampoo
+from .soap import SOAP
+from .orthograd import OrthoGrad, orthograd_
+from .sophia_h import SophiaH
+# from .curveball import CurveBall
+# from .spectral import SpectralPreconditioner