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

torchzero/modules/{momentum/ema.py → ops/higher_level.py}

@@ -5,39 +5,16 @@ from typing import Literal
 import torch
 
 from ...core import Target, Transform
-from ...utils import TensorList, NumberList, unpack_dicts, unpack_states
-from ..functional import debias, ema_, ema_sq_, sqrt_ema_sq_, centered_ema_sq_, sqrt_centered_ema_sq_, debias_second_momentum
-
-
-class EMA(Transform):
-    """Maintains an exponential moving average of update.
-
-    Args:
-        momentum (float, optional): momentum (beta). Defaults to 0.9.
-        dampening (float, optional): momentum dampening. Defaults to 0.
-        debiased (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
-        lerp (bool, optional): whether to use linear interpolation. Defaults to True.
-        ema_init (str, optional): initial values for the EMA, "zeros" or "update".
-        target (Target, optional): target to apply EMA to. Defaults to 'update'.
-    """
-    def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=True, ema_init: Literal['zeros', 'update'] = 'zeros', target: Target = 'update'):
-        defaults = dict(momentum=momentum,dampening=dampening,debiased=debiased,lerp=lerp,ema_init=ema_init)
-        super().__init__(defaults, uses_grad=False, target=target)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
-
-        debiased, lerp, ema_init = itemgetter('debiased','lerp','ema_init')(settings[0])
-
-        exp_avg = unpack_states(states, tensors, 'exp_avg',
-                                init=torch.zeros_like if ema_init=='zeros' else tensors, cls=TensorList)
-        momentum, dampening = unpack_dicts(settings, 'momentum','dampening', cls=NumberList)
-
-        exp_avg = ema_(TensorList(tensors), exp_avg_=exp_avg,beta=momentum,dampening=dampening,lerp=lerp)
-
-        if debiased: return debias(exp_avg, step=step, beta1=momentum, alpha=1, inplace=False)
-        else: return exp_avg.clone() # this has exp_avg storage so needs to be cloned
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+from ..functional import (
+    centered_ema_sq_,
+    debias,
+    debias_second_momentum,
+    ema_,
+    ema_sq_,
+    sqrt_centered_ema_sq_,
+    sqrt_ema_sq_,
+)
 
 
 class EMASquared(Transform):
@@ -55,7 +32,7 @@ class EMASquared(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):
         amsgrad, pow = itemgetter('amsgrad', 'pow')(self.settings[params[0]])
         beta = NumberList(s['beta'] for s in settings)
 
@@ -83,7 +60,7 @@ class SqrtEMASquared(Transform):
 
 
     @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['step'] = self.global_state.get('step', 0) + 1
 
         amsgrad, pow, debiased = itemgetter('amsgrad', 'pow', 'debiased')(settings[0])
@@ -123,7 +100,7 @@ class Debias(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['step'] = self.global_state.get('step', 0) + 1
 
         pow = settings[0]['pow']
@@ -145,7 +122,7 @@ class Debias2(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['step'] = self.global_state.get('step', 0) + 1
 
         pow = settings[0]['pow']
@@ -166,7 +143,7 @@ class CenteredEMASquared(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):
         amsgrad, pow = itemgetter('amsgrad', 'pow')(settings[0])
         beta = NumberList(s['beta'] for s in settings)
 
@@ -200,7 +177,7 @@ class CenteredSqrtEMASquared(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['step'] = self.global_state.get('step', 0) + 1
 
         amsgrad, pow, debiased = itemgetter('amsgrad', 'pow', 'debiased')(settings[0])

torchzero/modules/ops/multi.py

@@ -3,15 +3,15 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence
 from operator import itemgetter
-from typing import Any
+from typing import Any, Literal
 
 import torch
 
 from ...core import Chainable, Module, Target, Var, maybe_chain
-from ...utils import TensorList, tensorlist
+from ...utils import TensorList, tensorlist, Metrics
 
 
-class MultiOperation(Module, ABC):
+class MultiOperationBase(Module, ABC):
     """Base class for operations that use operands. 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)
@@ -51,14 +51,15 @@ class MultiOperation(Module, ABC):
 
 
 
-class SubModules(MultiOperation):
+class SubModules(MultiOperationBase):
+    """Calculates :code:`input - other`. :code:`input` and :code:`other` can be numbers or modules."""
     def __init__(self, input: Chainable | float, other: Chainable | float, alpha: float = 1):
         defaults = dict(alpha=alpha)
         super().__init__(defaults, input=input, other=other)
 
     @torch.no_grad
     def transform(self, var: Var, input: float | list[torch.Tensor], other: float | list[torch.Tensor]) -> list[torch.Tensor]:
-        alpha = self.settings[var.params[0]]['alpha']
+        alpha = self.defaults['alpha']
 
         if isinstance(input, (int,float)):
             assert isinstance(other, list)
@@ -68,10 +69,12 @@ class SubModules(MultiOperation):
         else: torch._foreach_sub_(input, other, alpha=alpha)
         return input
 
-class DivModules(MultiOperation):
-    def __init__(self, input: Chainable | float, other: Chainable | float):
+class DivModules(MultiOperationBase):
+    """Calculates :code:`input / other`. :code:`input` and :code:`other` can be numbers or modules."""
+    def __init__(self, input: Chainable | float, other: Chainable | float, other_first:bool=False):
         defaults = {}
-        super().__init__(defaults, input=input, other=other)
+        if other_first: super().__init__(defaults, other=other, input=input)
+        else: super().__init__(defaults, input=input, other=other)
 
     @torch.no_grad
     def transform(self, var: Var, input: float | list[torch.Tensor], other: float | list[torch.Tensor]) -> list[torch.Tensor]:
@@ -82,7 +85,9 @@ class DivModules(MultiOperation):
         torch._foreach_div_(input, other)
         return input
 
-class PowModules(MultiOperation):
+
+class PowModules(MultiOperationBase):
+    """Calculates :code:`input ** exponent`. :code:`input` and :code:`other` can be numbers or modules."""
     def __init__(self, input: Chainable | float, exponent: Chainable | float):
         defaults = {}
         super().__init__(defaults, input=input, exponent=exponent)
@@ -96,17 +101,22 @@ class PowModules(MultiOperation):
         torch._foreach_div_(input, exponent)
         return input
 
-class LerpModules(MultiOperation):
+class LerpModules(MultiOperationBase):
+    """Does a linear interpolation of :code:`input(tensors)` and :code:`end(tensors)` based on a scalar :code:`weight`.
+
+    The output is given by :code:`output = input(tensors) + weight * (end(tensors) - input(tensors))`
+    """
     def __init__(self, input: Chainable, end: Chainable, weight: float):
         defaults = dict(weight=weight)
         super().__init__(defaults, input=input, end=end)
 
     @torch.no_grad
     def transform(self, var: Var, input: list[torch.Tensor], end: list[torch.Tensor]) -> list[torch.Tensor]:
-        torch._foreach_lerp_(input, end, weight=self.settings[var.params[0]]['weight'])
+        torch._foreach_lerp_(input, end, weight=self.defaults['weight'])
         return input
 
-class ClipModules(MultiOperation):
+class ClipModules(MultiOperationBase):
+    """Calculates :code:`input(tensors).clip(min, max)`. :code:`min` and :code:`max` can be numbers or modules."""
     def __init__(self, input: Chainable, min: float | Chainable | None = None, max: float | Chainable | None = None):
         defaults = {}
         super().__init__(defaults, input=input, min=min, max=max)
@@ -116,22 +126,73 @@ class ClipModules(MultiOperation):
         return TensorList(input).clamp_(min=min, max=max)
 
 
-class GraftModules(MultiOperation):
-    def __init__(self, direction: Chainable, magnitude: Chainable, tensorwise:bool=True, ord:float=2, eps:float = 1e-6, strength:float=1):
+class GraftModules(MultiOperationBase):
+    """Outputs :code:`direction` output rescaled to have the same norm as :code:`magnitude` output.
+
+    Args:
+        direction (Chainable): module to use the direction from
+        magnitude (Chainable): module to use the magnitude from
+        tensorwise (bool, optional): whether to calculate norm per-tensor or globally. Defaults to True.
+        ord (float, optional): norm order. Defaults to 2.
+        eps (float, optional): clips denominator to be no less than this value. Defaults to 1e-6.
+        strength (float, optional): strength of grafting. Defaults to 1.
+
+    Example:
+        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)
+            )
+
+    Reference:
+        Agarwal, N., Anil, R., Hazan, E., Koren, T., & Zhang, C. (2020). Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11803. https://arxiv.org/pdf/2002.11803
+    """
+    def __init__(self, direction: Chainable, magnitude: Chainable, tensorwise:bool=True, ord:Metrics=2, eps:float = 1e-6, strength:float=1):
         defaults = dict(tensorwise=tensorwise, ord=ord, eps=eps, strength=strength)
         super().__init__(defaults, direction=direction, magnitude=magnitude)
 
     @torch.no_grad
     def transform(self, var, magnitude: list[torch.Tensor], direction:list[torch.Tensor]):
-        tensorwise, ord, eps, strength = itemgetter('tensorwise','ord','eps', 'strength')(self.settings[var.params[0]])
+        tensorwise, ord, eps, strength = itemgetter('tensorwise','ord','eps', 'strength')(self.defaults)
         return TensorList(direction).graft_(magnitude, tensorwise=tensorwise, ord=ord, eps=eps, strength=strength)
 
-
-class Where(MultiOperation):
-    def __init__(self, condition: Chainable, input: Chainable | float, other: Chainable | float):
-        super().__init__({}, condition=condition, input=input, other=other)
+class MultiplyByModuleNorm(MultiOperationBase):
+    """Outputs :code:`input` multiplied by norm of the :code:`norm` output."""
+    def __init__(self, input: Chainable, norm: Chainable, tensorwise:bool=True, ord:Metrics=2):
+        defaults = dict(tensorwise=tensorwise, ord=ord)
+        super().__init__(defaults, input=input, norm=norm)
 
     @torch.no_grad
-    def transform(self, var, condition: list[torch.Tensor], input: list[torch.Tensor] | float, other: list[torch.Tensor] | float):
-        return tensorlist.where(TensorList(condition).as_bool(), input, other)
+    def transform(self, var, input: list[torch.Tensor], norm:list[torch.Tensor]):
+        tensorwise, ord = itemgetter('tensorwise','ord')(self.defaults)
+        if tensorwise:
+            n = TensorList(norm).metric(ord)
+        else:
+            n = TensorList(norm).global_metric(ord)
+
+        torch._foreach_mul_(input, n)
+        return input
+
+class DivideByModuleNorm(MultiOperationBase):
+    """Outputs :code:`input` divided by norm of the :code:`norm` output."""
+    def __init__(self, input: Chainable, norm: Chainable, tensorwise:bool=True, ord:Metrics=2):
+        defaults = dict(tensorwise=tensorwise, ord=ord)
+        super().__init__(defaults, input=input, norm=norm)
 
+    @torch.no_grad
+    def transform(self, var, input: list[torch.Tensor], norm:list[torch.Tensor]):
+        tensorwise, ord = itemgetter('tensorwise','ord')(self.defaults)
+        if tensorwise:
+            n = TensorList(norm).metric(ord)
+        else:
+            n = TensorList(norm).global_metric(ord)
+
+        torch._foreach_div_(input, n)
+        return input

torchzero/modules/ops/reduce.py

@@ -8,7 +8,7 @@ import torch
 from ...core import Chainable, Module, Target, Var, maybe_chain
 
 
-class ReduceOperation(Module, ABC):
+class ReduceOperationBase(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)
@@ -46,7 +46,8 @@ class ReduceOperation(Module, ABC):
         var.update = transformed
         return var
 
-class Sum(ReduceOperation):
+class Sum(ReduceOperationBase):
+    """Outputs sum of :code:`inputs` that can be modules or numbers."""
     USE_MEAN = False
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
@@ -63,12 +64,14 @@ class Sum(ReduceOperation):
         return sum
 
 class Mean(Sum):
+    """Outputs a mean of :code:`inputs` that can be modules or numbers."""
     USE_MEAN = True
 
 
-class WeightedSum(ReduceOperation):
+class WeightedSum(ReduceOperationBase):
     USE_MEAN = False
     def __init__(self, *inputs: Chainable | float, weights: Iterable[float]):
+        """Outputs a weighted sum of :code:`inputs` that can be modules or numbers."""
         weights = list(weights)
         if len(inputs) != len(weights):
             raise ValueError(f'Number of inputs {len(inputs)} must match number of weights {len(weights)}')
@@ -78,7 +81,7 @@ class WeightedSum(ReduceOperation):
     @torch.no_grad
     def transform(self, var: Var, *inputs: float | list[torch.Tensor]) -> list[torch.Tensor]:
         sorted_inputs = sorted(inputs, key=lambda x: isinstance(x, float))
-        weights = self.settings[var.params[0]]['weights']
+        weights = self.defaults['weights']
         sum = cast(list, sorted_inputs[0])
         torch._foreach_mul_(sum, weights[0])
         if len(sorted_inputs) > 1:
@@ -91,9 +94,11 @@ class WeightedSum(ReduceOperation):
 
 
 class WeightedMean(WeightedSum):
+    """Outputs weighted mean of :code:`inputs` that can be modules or numbers."""
     USE_MEAN = True
 
-class Median(ReduceOperation):
+class Median(ReduceOperationBase):
+    """Outputs median of :code:`inputs` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
@@ -106,7 +111,8 @@ class Median(ReduceOperation):
             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):
+class Prod(ReduceOperationBase):
+    """Outputs product of :code:`inputs` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
@@ -120,7 +126,8 @@ class Prod(ReduceOperation):
 
         return prod
 
-class MaximumModules(ReduceOperation):
+class MaximumModules(ReduceOperationBase):
+    """Outputs elementwise maximum of :code:`inputs` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 
@@ -134,7 +141,8 @@ class MaximumModules(ReduceOperation):
 
         return maximum
 
-class MinimumModules(ReduceOperation):
+class MinimumModules(ReduceOperationBase):
+    """Outputs elementwise minimum of :code:`inputs` that can be modules or numbers."""
     def __init__(self, *inputs: Chainable | float):
         super().__init__({}, *inputs)
 

torchzero/modules/ops/unary.py

@@ -6,76 +6,92 @@ from ...core import TensorwiseTransform, Target, Transform
 from ...utils import TensorList, unpack_dicts,unpack_states
 
 class UnaryLambda(Transform):
+    """Applies :code:`fn` to input tensors.
+
+    :code:`fn` must accept and return a list of tensors.
+    """
     def __init__(self, fn, target: "Target" = 'update'):
         defaults = dict(fn=fn)
         super().__init__(defaults=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):
         return settings[0]['fn'](tensors)
 
 class UnaryParameterwiseLambda(TensorwiseTransform):
+    """Applies :code:`fn` to each input tensor.
+
+    :code:`fn` must accept and return a tensor.
+    """
     def __init__(self, fn, target: "Target" = 'update'):
         defaults = dict(fn=fn)
         super().__init__(uses_grad=False, defaults=defaults, target=target)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
-        return settings['fn'](tensor)
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        return setting['fn'](tensor)
 
 class CustomUnaryOperation(Transform):
+    """Applies :code:`getattr(tensor, name)` to each tensor
+    """
     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 apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         return getattr(tensors, settings[0]['name'])()
 
 
 class Abs(Transform):
+    """Returns :code:`abs(input)`"""
     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):
         torch._foreach_abs_(tensors)
         return tensors
 
 class Sign(Transform):
+    """Returns :code:`sign(input)`"""
     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):
         torch._foreach_sign_(tensors)
         return tensors
 
 class Exp(Transform):
+    """Returns :code:`exp(input)`"""
     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):
         torch._foreach_exp_(tensors)
         return tensors
 
 class Sqrt(Transform):
+    """Returns :code:`sqrt(input)`"""
     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):
         torch._foreach_sqrt_(tensors)
         return tensors
 
 class Reciprocal(Transform):
+    """Returns :code:`1 / input`"""
     def __init__(self, eps = 0, target: "Target" = 'update'):
         defaults = dict(eps = eps)
         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):
         eps = [s['eps'] for s in settings]
         if any(e != 0 for e in eps): torch._foreach_add_(tensors, eps)
         torch._foreach_reciprocal_(tensors)
         return tensors
 
 class Negate(Transform):
+    """Returns :code:`- input`"""
     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):
         torch._foreach_neg_(tensors)
         return tensors
 
@@ -97,18 +113,18 @@ class NanToNum(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):
         nan, posinf, neginf = unpack_dicts(settings, 'nan', 'posinf', 'neginf')
         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"""
+    """Rescales input to :code`(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 apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         min, max = unpack_dicts(settings, 'min','max')
         tensorwise = settings[0]['tensorwise']
         dim = None if tensorwise else 'global'

torchzero/modules/ops/utility.py

@@ -4,38 +4,37 @@ import torch
 
 from ...core import Module, Target, Transform
 from ...utils.tensorlist import Distributions, TensorList
+from ...utils.linalg.linear_operator import ScaledIdentity
 
-
-class Clone(Transform):
-    def __init__(self): super().__init__({}, uses_grad=False)
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings): return [t.clone() for t in tensors]
-
-class Grad(Module):
+class Clone(Module):
+    """Clones input. May be useful to store some intermediate result and make sure it doesn't get affected by in-place operations"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
     def step(self, var):
-        var.update = [g.clone() for g in var.get_grad()]
+        var.update = [u.clone() for u in var.get_update()]
         return var
 
-class Params(Module):
+class Grad(Module):
+    """Outputs the gradient"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
     def step(self, var):
-        var.update = [p.clone() for p in var.params]
+        var.update = [g.clone() for g in var.get_grad()]
         return var
 
-class Update(Module):
+class Params(Module):
+    """Outputs parameters"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
     def step(self, var):
-        var.update = [u.clone() for u in var.get_update()]
+        var.update = [p.clone() for p in var.params]
         return var
 
 class Zeros(Module):
+    """Outputs zeros"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
@@ -44,6 +43,7 @@ class Zeros(Module):
         return var
 
 class Ones(Module):
+    """Outputs ones"""
     def __init__(self):
         super().__init__({})
     @torch.no_grad
@@ -52,6 +52,7 @@ class Ones(Module):
         return var
 
 class Fill(Module):
+    """Outputs tensors filled with :code:`value`"""
     def __init__(self, value: float):
         defaults = dict(value=value)
         super().__init__(defaults)
@@ -62,18 +63,20 @@ class Fill(Module):
         return var
 
 class RandomSample(Module):
-    def __init__(self, eps: float = 1, distribution: Distributions = 'normal'):
-        defaults = dict(eps=eps, distribution=distribution)
+    """Outputs tensors filled with random numbers from distribution depending on value of :code:`distribution`."""
+    def __init__(self, distribution: Distributions = 'normal', variance:float | None = None):
+        defaults = dict(distribution=distribution, variance=variance)
         super().__init__(defaults)
 
     @torch.no_grad
     def step(self, var):
-        var.update = TensorList(var.params).sample_like(
-            eps=[self.settings[p]['eps'] for p in var.params], distribution=self.settings[var.params[0]]['distribution']
-        )
+        distribution = self.defaults['distribution']
+        variance = self.get_settings(var.params, 'variance')
+        var.update = TensorList(var.params).sample_like(distribution=distribution, variance=variance)
         return var
 
 class Randn(Module):
+    """Outputs tensors filled with random numbers from a normal distribution with mean 0 and variance 1."""
     def __init__(self):
         super().__init__({})
 
@@ -83,6 +86,7 @@ class Randn(Module):
         return var
 
 class Uniform(Module):
+    """Outputs tensors filled with random numbers from uniform distribution between :code:`low` and :code:`high`."""
     def __init__(self, low: float, high: float):
         defaults = dict(low=low, high=high)
         super().__init__(defaults)
@@ -94,19 +98,27 @@ class Uniform(Module):
         return var
 
 class GradToNone(Module):
+    """Sets :code:`grad` attribute to None on :code:`var`."""
     def __init__(self): super().__init__()
     def step(self, var):
         var.grad = None
         return var
 
 class UpdateToNone(Module):
+    """Sets :code:`update` attribute to None on :code:`var`."""
     def __init__(self): super().__init__()
     def step(self, var):
         var.update = None
         return var
 
 class Identity(Module):
+    """Identity operator that is argument-insensitive. This also can be used as identity hessian for trust region methods."""
     def __init__(self, *args, **kwargs): super().__init__()
     def step(self, var): return var
+    def get_H(self, var):
+        n = sum(p.numel() for p in var.params)
+        p = var.params[0]
+        return ScaledIdentity(shape=(n,n), device=p.device, dtype=p.dtype)
 
-NoOp = Identity
+Noop = Identity
+"""A placeholder identity operator that is argument-insensitive."""

torchzero/modules/projections/__init__.py

@@ -1,5 +1,3 @@
-from .projection import Projection
-from .fft import FFTProjection
-from .structural import VectorProjection, TensorizeProjection, BlockPartition, TensorNormsProjection
-
+from .projection import ProjectionBase, VectorProjection, ScalarProjection
+from .cast import To, ViewAsReal
 # from .galore import GaLore