torchzero 0.0.1__py3-none-any.whl

torchzero/modules/operations/multi.py

@@ -0,0 +1,298 @@
+from collections.abc import Iterable
+import torch
+
+from ...core import OptimizerModule
+
+_Value = int | float | OptimizerModule | Iterable[OptimizerModule]
+
+class Add(OptimizerModule):
+    """add `value` to update. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
+    def __init__(self, value: _Value):
+        super().__init__({})
+
+        if not isinstance(value, (int, float)):
+            self._set_child_('value', value)
+
+        self.value = value
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if isinstance(self.value, (int, float)):
+            return ascent.add_(self.value)
+
+        state_copy = state.copy(clone_ascent = True)
+        v = self.children['value'].return_ascent(state_copy)
+        return ascent.add_(v)
+
+
+class Sub(OptimizerModule):
+    """subtracts `value` from update. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
+    def __init__(self, subtrahend: _Value):
+        super().__init__({})
+
+        if not isinstance(subtrahend, (int, float)):
+            self._set_child_('subtrahend', subtrahend)
+
+        self.subtrahend = subtrahend
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if isinstance(self.subtrahend, (int, float)):
+            return ascent.sub_(self.subtrahend)
+
+        state_copy = state.copy(clone_ascent = True)
+        subtrahend = self.children['subtrahend'].return_ascent(state_copy)
+        return ascent.sub_(subtrahend)
+
+class RSub(OptimizerModule):
+    """subtracts update from `value`. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
+    def __init__(self, minuend: _Value):
+        super().__init__({})
+
+        if not isinstance(minuend, (int, float)):
+            self._set_child_('minuend', minuend)
+
+        self.minuend = minuend
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if isinstance(self.minuend, (int, float)):
+            return ascent.sub_(self.minuend).neg_()
+
+        state_copy = state.copy(clone_ascent = True)
+        minuend = self.children['minuend'].return_ascent(state_copy)
+        return ascent.sub_(minuend).neg_()
+
+class Subtract(OptimizerModule):
+    """Calculates `minuend - subtrahend`"""
+    def __init__(
+        self,
+        minuend: OptimizerModule | Iterable[OptimizerModule],
+        subtrahend: OptimizerModule | Iterable[OptimizerModule],
+    ):
+        super().__init__({})
+        self._set_child_('minuend', minuend)
+        self._set_child_('subtrahend', subtrahend)
+
+    @torch.no_grad
+    def step(self, state):
+        state_copy = state.copy(clone_ascent = True)
+        minuend = self.children['minuend'].return_ascent(state_copy)
+        state.update_attrs_(state_copy)
+        subtrahend = self.children['subtrahend'].return_ascent(state)
+
+        state.ascent = minuend.sub_(subtrahend)
+        return self._update_params_or_step_with_next(state)
+
+class Mul(OptimizerModule):
+    """multiplies update by `value`. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
+    def __init__(self, value: _Value):
+        super().__init__({})
+
+        if not isinstance(value, (int, float)):
+            self._set_child_('value', value)
+
+        self.value = value
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if isinstance(self.value, (int, float)):
+            return ascent.mul_(self.value)
+
+        state_copy = state.copy(clone_ascent = True)
+        v = self.children['value'].return_ascent(state_copy)
+        return ascent.mul_(v)
+
+
+class Div(OptimizerModule):
+    """divides update by `value`. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
+    def __init__(self, denominator: _Value):
+        super().__init__({})
+
+        if not isinstance(denominator, (int, float)):
+            self._set_child_('denominator', denominator)
+
+        self.denominator = denominator
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if isinstance(self.denominator, (int, float)):
+            return ascent.div_(self.denominator)
+
+        state_copy = state.copy(clone_ascent = True)
+        denominator = self.children['denominator'].return_ascent(state_copy)
+        return ascent.div_(denominator)
+
+class RDiv(OptimizerModule):
+    """`value` by update. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
+    def __init__(self, numerator: _Value):
+        super().__init__({})
+
+        if not isinstance(numerator, (int, float)):
+            self._set_child_('numerator', numerator)
+
+        self.numerator = numerator
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if isinstance(self.numerator, (int, float)):
+            return ascent.reciprocal_().mul_(self.numerator)
+
+        state_copy = state.copy(clone_ascent = True)
+        numerator = self.children['numerator'].return_ascent(state_copy)
+        return ascent.reciprocal_().mul_(numerator)
+
+class Divide(OptimizerModule):
+    """calculates *numerator / denominator*"""
+    def __init__(
+        self,
+        numerator: OptimizerModule | Iterable[OptimizerModule],
+        denominator: OptimizerModule | Iterable[OptimizerModule],
+    ):
+        super().__init__({})
+        self._set_child_('numerator', numerator)
+        self._set_child_('denominator', denominator)
+
+    @torch.no_grad
+    def step(self, state):
+        state_copy = state.copy(clone_ascent = True)
+        numerator = self.children['numerator'].return_ascent(state_copy)
+        state.update_attrs_(state_copy)
+        denominator = self.children['denominator'].return_ascent(state)
+
+        state.ascent = numerator.div_(denominator)
+        return self._update_params_or_step_with_next(state)
+
+
+class Pow(OptimizerModule):
+    """takes ascent to the power of `value`. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
+    def __init__(self, power: _Value):
+        super().__init__({})
+
+        if not isinstance(power, (int, float)):
+            self._set_child_('power', power)
+
+        self.power = power
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if isinstance(self.power, (int, float)):
+            return ascent.pow_(self.power)
+
+        state_copy = state.copy(clone_ascent = True)
+        power = self.children['power'].return_ascent(state_copy)
+        return ascent.pow_(power)
+
+class RPow(OptimizerModule):
+    """takes `value` to the power of ascent. `value` can be a scalar, an OptimizerModule or sequence of OptimizerModules"""
+    def __init__(self, base: _Value):
+        super().__init__({})
+
+        if not isinstance(base, (int, float)):
+            self._set_child_('base', base)
+
+        self.base = base
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if isinstance(self.base, (int, float)):
+            return self.base ** ascent
+
+        state_copy = state.copy(clone_ascent = True)
+        base = self.children['base'].return_ascent(state_copy)
+        return base.pow_(ascent)
+
+class Power(OptimizerModule):
+    """calculates *base ^ power*"""
+    def __init__(
+        self,
+        base: OptimizerModule | Iterable[OptimizerModule],
+        power: OptimizerModule | Iterable[OptimizerModule],
+    ):
+        super().__init__({})
+        self._set_child_('base', base)
+        self._set_child_('power', power)
+
+    @torch.no_grad
+    def step(self, state):
+        state_copy = state.copy(clone_ascent = True)
+        base = self.children['base'].return_ascent(state_copy)
+        state.update_attrs_(state_copy)
+        power = self.children['power'].return_ascent(state)
+
+        state.ascent = base.pow_(power)
+        return self._update_params_or_step_with_next(state)
+
+
+class Lerp(OptimizerModule):
+    """Linear interpolation between update and `end` based on scalar `weight`.
+
+    `out = update + weight * (end - update)`"""
+    def __init__(self, end: OptimizerModule | Iterable[OptimizerModule], weight: float):
+        super().__init__({})
+
+        self._set_child_('end', end)
+        self.weight = weight
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+
+        state_copy = state.copy(clone_ascent = True)
+        end = self.children['end'].return_ascent(state_copy)
+        return ascent.lerp_(end, self.weight)
+
+
+class Interpolate(OptimizerModule):
+    """Does a linear interpolation of two module's updates - `start` (given by input), and `end`, based on a scalar
+    `weight`.
+
+    `out = input + weight * (end - input)`"""
+    def __init__(
+        self,
+        input: OptimizerModule | Iterable[OptimizerModule],
+        end: OptimizerModule | Iterable[OptimizerModule],
+        weight: float,
+    ):
+        super().__init__({})
+        self._set_child_('input', input)
+        self._set_child_('end', end)
+        self.weight = weight
+
+    @torch.no_grad
+    def step(self, state):
+        state_copy = state.copy(clone_ascent = True)
+        input = self.children['input'].return_ascent(state_copy)
+        state.update_attrs_(state_copy)
+        end = self.children['end'].return_ascent(state)
+
+        state.ascent = input.lerp_(end, weight = self.weight)
+
+        return self._update_params_or_step_with_next(state)
+
+class AddMagnitude(OptimizerModule):
+    """Add `value` multiplied by sign of the ascent, i.e. this adds `value` to the magnitude of the update.
+
+    Args:
+        value (Value): value to add to magnitude, either a float or an OptimizerModule.
+        add_to_zero (bool, optional):
+            if True, adds `value` to 0s. Otherwise, zeros remain zero.
+            Only has effect if value is a float. Defaults to True.
+    """
+    def __init__(self, value: _Value, add_to_zero=True):
+        super().__init__({})
+
+        if not isinstance(value, (int, float)):
+            self._set_child_('value', value)
+
+        self.value = value
+        self.add_to_zero = add_to_zero
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if isinstance(self.value, (int, float)):
+            if self.add_to_zero: return ascent.add_(ascent.clamp_magnitude(min=1).sign_().mul_(self.value))
+            return ascent.add_(ascent.sign_().mul_(self.value))
+
+        state_copy = state.copy(clone_ascent = True)
+        v = self.children['value'].return_ascent(state_copy)
+        return ascent.add_(v.abs_().mul_(ascent.sign()))

torchzero/modules/operations/reduction.py

@@ -0,0 +1,134 @@
+from collections.abc import Callable, Iterable
+import numpy as np
+import torch
+
+from ...core import OptimizerModule
+
+_Value = int | float | OptimizerModule | Iterable[OptimizerModule]
+
+
+class Sum(OptimizerModule):
+    """calculates sum of multiple updates.
+
+    Args:
+        *modules:
+            either OptimizerModules or iterables of OptimizerModules to chain. Scalars are also allowed."""
+    def __init__(
+        self,
+        *modules: _Value,
+    ):
+        super().__init__({})
+
+        scalars = [i for i in modules if isinstance(i, (int,float))]
+        self.scalar = sum(scalars) if len(scalars) > 0 else None
+
+        for i,module in enumerate(i for i in modules if not isinstance(i, (int, float))):
+            self._set_child_(i, module)
+
+    @torch.no_grad
+    def step(self, state):
+        if len(self.children) == 1:
+            state.ascent = self.children[0].return_ascent(state)
+            if self.scalar is not None: state.ascent += self.scalar
+            return self._update_params_or_step_with_next(state)
+
+        sum = None
+        for i, c in sorted(self.children.items(), key=lambda x: x[0]):
+            if i == len(self.children) - 1: cur_state = state
+            else: cur_state = state.copy(clone_ascent = True)
+
+            if sum is None: sum = c.return_ascent(cur_state)
+            else: sum += c.return_ascent(cur_state)
+
+            if i != len(self.children) - 1: state.update_attrs_(cur_state)
+
+        assert sum is not None
+        if self.scalar is not None: sum += self.scalar
+        state.ascent = sum
+        return self._update_params_or_step_with_next(state)
+
+class Mean(OptimizerModule):
+    """calculates mean of multiple updates.
+
+    Args:
+        *modules:
+            either OptimizerModules or iterables of OptimizerModules to chain. Scalars are also allowed."""
+
+    def __init__(
+        self,
+        *modules: _Value,
+    ):
+        super().__init__({})
+
+        scalars = [i for i in modules if isinstance(i, (int,float))]
+        self.scalar = sum(scalars) if len(scalars) > 0 else None
+
+        self.n_values = len(modules)
+
+        for i,module in enumerate(i for i in modules if not isinstance(i, (int, float))):
+            self._set_child_(i, module)
+
+    @torch.no_grad
+    def step(self, state):
+        if len(self.children) == 1:
+            state.ascent = self.children[0].return_ascent(state)
+            if self.scalar is not None: state.ascent += self.scalar
+            if self.n_values > 1: state.ascent /= self.n_values
+            return self._update_params_or_step_with_next(state)
+
+        sum = None
+        for i, c in sorted(self.children.items(), key=lambda x: x[0]):
+            if i == len(self.children) - 1: cur_state = state
+            else: cur_state = state.copy(clone_ascent = True)
+
+            if sum is None: sum = c.return_ascent(cur_state)
+            else: sum += c.return_ascent(cur_state)
+
+            if i != len(self.children) - 1: state.update_attrs_(cur_state)
+
+        assert sum is not None
+        if self.scalar is not None: sum += self.scalar
+        if self.n_values > 1: sum /= self.n_values
+        state.ascent = sum
+        return self._update_params_or_step_with_next(state)
+
+class Product(OptimizerModule):
+    """calculates product of multiple updates.
+
+    Args:
+        *modules:
+            either OptimizerModules or iterables of OptimizerModules to chain. Scalars are also allowed."""
+
+    def __init__(
+        self,
+        *modules: _Value,
+    ):
+        super().__init__({})
+
+        scalars = [i for i in modules if isinstance(i, (int,float))]
+        self.scalar = np.prod(scalars).item() if len(scalars) > 0 else None
+
+        for i,module in enumerate(i for i in modules if not isinstance(i, (int, float))):
+            self._set_child_(i, module)
+
+    @torch.no_grad
+    def step(self, state):
+        if len(self.children) == 1:
+            state.ascent = self.children[0].return_ascent(state)
+            if self.scalar is not None: state.ascent *= self.scalar
+            return self._update_params_or_step_with_next(state)
+
+        prod = None
+        for i, c in sorted(self.children.items(), key=lambda x: x[0]):
+            if i == len(self.children) - 1: cur_state = state
+            else: cur_state = state.copy(clone_ascent = True)
+
+            if prod is None: prod = c.return_ascent(cur_state)
+            else: prod *= c.return_ascent(cur_state)
+
+            if i != len(self.children) - 1: state.update_attrs_(cur_state)
+
+        assert prod is not None
+        if self.scalar is not None: prod *= self.scalar
+        state.ascent = prod
+        return self._update_params_or_step_with_next(state)

torchzero/modules/operations/singular.py

@@ -0,0 +1,113 @@
+from collections.abc import Iterable
+from operator import methodcaller
+
+import torch
+
+from ...core import OptimizerModule
+from ...tensorlist import TensorList
+
+
+class Operation(OptimizerModule):
+    """Applies an operation to the ascent, supported operations:
+
+    `abs`, `sign`, `sin`, `cos`, `tan`, `asin`, `acos`, `atan`, `sinh`, `cosh`,
+    `tanh`, `log`, `log1p`, `log2`, `log10`, `erf`, `erfc`, `exp`, `neg`, `reciprocal`,
+    `copy`, `zero`, `sqrt`, `floor`, `ceil`, `round`."""
+    def __init__(self, operation: str):
+        super().__init__({})
+        self.operation = methodcaller(f'{operation}_')
+
+    @torch.no_grad
+    def _update(self, state, ascent): return self.operation(ascent)
+
+class Reciprocal(OptimizerModule):
+    """*1 / update*"""
+    def __init__(self,):
+        super().__init__({})
+
+    @torch.no_grad()
+    def _update(self, state, ascent): return ascent.reciprocal_()
+
+class Negate(OptimizerModule):
+    """minus update"""
+    def __init__(self,):
+        super().__init__({})
+
+    @torch.no_grad()
+    def _update(self, state, ascent): return ascent.neg_()
+
+
+def sign_grad_(params: Iterable[torch.Tensor]):
+    """Apply sign function to gradients of an iterable of parameters.
+
+    Args:
+        params (abc.Iterable[torch.Tensor]): an iterable of Tensors or a single Tensor.
+    """
+    TensorList(params).get_existing_grads().sign_()
+
+class Sign(OptimizerModule):
+    """applies sign function to the update"""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def _update(self, state, ascent): return ascent.sign_()
+
+class Abs(OptimizerModule):
+    """takes absolute values of the update."""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def _update(self, state, ascent): return ascent.abs_()
+
+class Sin(OptimizerModule):
+    """applies sin function to the ascent"""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def _update(self, state, ascent): return ascent.sin_()
+
+class Cos(OptimizerModule):
+    """applies cos function to the ascent"""
+    def __init__(self):
+        super().__init__({})
+
+    @torch.no_grad
+    def _update(self, state, ascent): return ascent.cos_()
+
+
+class NanToNum(OptimizerModule):
+    """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):
+        super().__init__({})
+        self.nan = nan
+        self.posinf = posinf
+        self.neginf = neginf
+
+    @torch.no_grad()
+    def _update(self, state, ascent): return ascent.nan_to_num_(self.nan, self.posinf, self.neginf)
+
+
+class MagnitudePower(OptimizerModule):
+    """Raises update to the `value` power, but preserves the sign when the power is odd."""
+    def __init__(self, value: int | float):
+        super().__init__({})
+        self.value = value
+
+    @torch.no_grad()
+    def _update(self, state, ascent):
+        if self.value % 2 == 1: return ascent.pow_(self.value)
+        return ascent.abs().pow_(self.value) * ascent.sign()
+

torchzero/modules/optimizers/__init__.py

@@ -0,0 +1,10 @@
+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

torchzero/modules/optimizers/adagrad.py

@@ -0,0 +1,49 @@
+from collections import abc
+
+import torch
+
+from ...tensorlist import TensorList
+from ...core import OptimizerModule
+
+def _adagrad_step_(ascent: TensorList, grad_sum: TensorList, alpha: TensorList, lr_decay: TensorList, eps: TensorList, step: int):
+    clr = alpha / (1 + step * lr_decay)
+    grad_sum.addcmul_(ascent, ascent)
+    return ascent.div_(grad_sum.sqrt().add_(eps)).mul_(clr)
+
+class Adagrad(OptimizerModule):
+    """
+    Divides ascent direction by mean square root of the sum of all past ascent directions.
+
+    Exactly matches `torch.optim.Adagrad`.
+
+    Args:
+        lr_decay (float, optional): learning rate decay. Defaults to 0.
+        initial_accumulator_value (float, optional): initial value of the sum of squares of gradients. Defaults to 0.
+        eps (float, optional): term added to the denominator to improve numerical stability. Defaults to 1e-10.
+        alpha (float, optional): learning rate. Defaults to 1.
+
+    reference
+        https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf
+    """
+    def __init__(self, lr_decay: float = 0, initial_accumulator_value: float = 0, eps: float = 1e-10, alpha: float = 1):
+        defaults = dict(alpha = alpha, lr_decay = lr_decay, initial_accumulator_value=initial_accumulator_value, eps = eps)
+        super().__init__(defaults)
+        self.cur_step = 0
+
+    @torch.no_grad
+    def _update(self, state, ascent):
+        settings = self.get_all_group_keys()
+        if self.cur_step == 0: init = ascent.full_like(settings['initial_accumulator_value'])
+        else: init = None
+        grad_sum = self.get_state_key('grad_sum', init = init) # type:ignore
+
+        updated_direction = _adagrad_step_(
+            ascent=ascent,
+            grad_sum=grad_sum,
+            alpha=settings["alpha"],
+            eps=settings["eps"],
+            lr_decay=settings["lr_decay"],
+            step=self.cur_step,
+        )
+        self.cur_step += 1
+        return updated_direction