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

torchzero/modules/experimental/experimental.py

@@ -1,294 +0,0 @@
-from collections import abc
-
-import torch
-
-from ...tensorlist import TensorList, where, Distributions
-from ...core import OptimizerModule
-from ...utils.derivatives import jacobian
-
-def _bool_ones_like(x):
-    return torch.ones_like(x, dtype=torch.bool)
-
-
-class MinibatchRprop(OptimizerModule):
-    """
-    for experiments, unlikely to work well on most problems.
-
-    explanation: does 2 steps per batch, applies rprop rule on the second step.
-    """
-    def __init__(
-        self,
-        nplus: float = 1.2,
-        nminus: float = 0.5,
-        lb: float | None = 1e-6,
-        ub: float | None = 50,
-        backtrack=True,
-        next_mode = 'continue',
-        increase_mul = 0.5,
-        alpha: float = 1,
-    ):
-        defaults = dict(nplus = nplus, nminus = nminus, alpha = alpha, lb = lb, ub = ub, increase_mul=increase_mul)
-        super().__init__(defaults)
-        self.current_step = 0
-        self.backtrack = backtrack
-
-        self.next_mode = next_mode
-
-    @torch.no_grad
-    def step(self, vars):
-        if vars.closure is None: raise ValueError("Minibatch Rprop requires closure")
-        if vars.ascent is not None: raise ValueError("Minibatch Rprop must be the first module.")
-        params = self.get_params()
-
-        nplus, nminus, lb, ub = self.get_group_keys('nplus', 'nminus', 'lb', 'ub')
-        allowed, magnitudes = self.get_state_keys(
-            'allowed', 'magnitudes',
-            inits = [_bool_ones_like, torch.zeros_like],
-            params=params
-        )
-
-        g1_sign = vars.maybe_compute_grad_(params).sign() # no inplace to not modify grads
-        # initialize on 1st iteration
-        if self.current_step == 0:
-            magnitudes.fill_(self.get_group_key('alpha')).clamp_(lb, ub)
-            # ascent = magnitudes * g1_sign
-            # self.current_step += 1
-            # return ascent
-
-        # first step
-        ascent = g1_sign.mul_(magnitudes).mul_(allowed)
-        params -= ascent
-        with torch.enable_grad(): vars.fx0_approx = vars.closure()
-        f0 = vars.fx0; f1 = vars.fx0_approx
-        assert f0 is not None and f1 is not None
-
-        # if loss increased, reduce all lrs and undo the update
-        if f1 > f0:
-            increase_mul = self.get_group_key('increase_mul')
-            magnitudes.mul_(increase_mul).clamp_(lb, ub)
-            params += ascent
-            self.current_step += 1
-            return f0
-
-        # on `continue` we move to params after 1st update
-        # therefore state must be updated to have all attributes after 1st update
-        if self.next_mode == 'continue':
-            vars.fx0 = vars.fx0_approx
-            vars.grad = params.ensure_grad_().grad
-            sign = vars.grad.sign()
-
-        else:
-            sign = params.ensure_grad_().grad.sign_() # can use in-place as this is not fx0 grad
-
-        # compare 1st and 2nd gradients via rprop rule
-        prev = ascent
-        mul = sign * prev # prev is already multiuplied by `allowed`
-
-        sign_changed = mul < 0
-        sign_same = mul > 0
-        zeroes = mul == 0
-
-        mul.fill_(1)
-        mul.masked_fill_(sign_changed, nminus)
-        mul.masked_fill_(sign_same, nplus)
-
-        # multiply magnitudes based on sign change and clamp to bounds
-        magnitudes.mul_(mul).clamp_(lb, ub)
-
-        # revert update if sign changed
-        if self.backtrack:
-            ascent2 = sign.mul_(magnitudes)
-            ascent2.masked_set_(sign_changed, prev.neg_())
-        else:
-            ascent2 = sign.mul_(magnitudes * ~sign_changed)
-
-        # update allowed to only have weights where last update wasn't reverted
-        allowed.set_(sign_same | zeroes)
-
-        self.current_step += 1
-
-        # update params or step
-        if self.next_mode == 'continue' or (self.next_mode == 'add' and self.next_module is None):
-            vars.ascent = ascent2
-            return self._update_params_or_step_with_next(vars, params)
-
-        if self.next_mode == 'add':
-            # undo 1st step
-            params += ascent
-            vars.ascent = ascent + ascent2
-            return self._update_params_or_step_with_next(vars, params)
-
-        if self.next_mode == 'undo':
-            params += ascent
-            vars.ascent = ascent2
-            return self._update_params_or_step_with_next(vars, params)
-
-        raise ValueError(f'invalid next_mode: {self.next_mode}')
-
-
-
-class GradMin(OptimizerModule):
-    """
-    for experiments, unlikely to work well on most problems.
-
-    explanation: calculate grads wrt sum of grads + loss.
-    """
-    def __init__(self, loss_term: float = 1, square=False, maximize_grad = False, create_graph = False):
-        super().__init__(dict(loss_term=loss_term))
-        self.square = square
-        self.maximize_grad = maximize_grad
-        self.create_graph = create_graph
-
-    @torch.no_grad
-    def step(self, vars):
-        if vars.closure is None: raise ValueError()
-        if vars.ascent is not None:
-            raise ValueError("GradMin doesn't accept ascent_direction")
-
-        params = self.get_params()
-        loss_term = self.get_group_key('loss_term')
-
-        self.zero_grad()
-        with torch.enable_grad():
-            vars.fx0 = vars.closure(False)
-            grads = jacobian([vars.fx0], params, create_graph=True, batched=False) # type:ignore
-            grads = TensorList(grads).squeeze_(0)
-            if self.square:
-                grads = grads ** 2
-            else:
-                grads = grads.abs()
-
-            if self.maximize_grad: grads: TensorList = grads - (vars.fx0 * loss_term) # type:ignore
-            else: grads = grads + (vars.fx0 * loss_term)
-            grad_mean = torch.sum(torch.stack(grads.sum())) / grads.total_numel()
-
-            if self.create_graph: grad_mean.backward(create_graph=True)
-            else: grad_mean.backward(retain_graph=False)
-
-        if self.maximize_grad: vars.grad = params.ensure_grad_().grad.neg_()
-        else: vars.grad = params.ensure_grad_().grad
-
-        vars.maybe_use_grad_(params)
-        return self._update_params_or_step_with_next(vars)
-
-
-class HVPDiagNewton(OptimizerModule):
-    """
-    for experiments, unlikely to work well on most problems.
-
-    explanation: may or may not approximate newton step if hessian is diagonal with 2 backward passes. Probably not.
-    """
-    def __init__(self, eps=1e-3):
-        super().__init__(dict(eps=eps))
-
-    @torch.no_grad
-    def step(self, vars):
-        if vars.closure is None: raise ValueError()
-        if vars.ascent is not None:
-            raise ValueError("HVPDiagNewton doesn't accept ascent_direction")
-
-        params = self.get_params()
-        eps = self.get_group_key('eps')
-        grad_fx0 = vars.maybe_compute_grad_(params).clone()
-        vars.grad = grad_fx0 # set state grad to the cloned version, since it will be overwritten
-
-        params += grad_fx0 * eps
-        with torch.enable_grad(): _ = vars.closure()
-
-        params -= grad_fx0 * eps
-
-        newton = grad_fx0 * ((grad_fx0 * eps) / (params.grad - grad_fx0))
-        newton.nan_to_num_(0,0,0)
-
-        vars.ascent = newton
-        return self._update_params_or_step_with_next(vars)
-
-
-
-class ReduceOutwardLR(OptimizerModule):
-    """
-    When update sign matches weight sign, the learning rate for that weight is multiplied by `mul`.
-
-    This means updates that move weights towards zero have higher learning rates.
-    """
-    def __init__(self, mul = 0.5, use_grad=False, invert=False):
-        defaults = dict(mul = mul)
-        super().__init__(defaults)
-
-        self.use_grad = use_grad
-        self.invert = invert
-
-    @torch.no_grad
-    def _update(self, vars, ascent):
-        params = self.get_params()
-        mul = self.get_group_key('mul')
-
-        if self.use_grad: cur = vars.maybe_compute_grad_(params)
-        else: cur = ascent
-
-        # mask of weights where sign matches with update sign (minus ascent sign), multiplied by `mul`.
-        if self.invert: mask = (params * cur) > 0
-        else: mask = (params * cur) < 0
-        ascent.masked_set_(mask, ascent*mul)
-
-        return ascent
-
-class NoiseSign(OptimizerModule):
-    """uses random vector with ascent sign"""
-    def __init__(self, distribution:Distributions = 'normal', alpha = 1):
-        super().__init__({})
-        self.alpha = alpha
-        self.distribution:Distributions = distribution
-
-
-    def _update(self, vars, ascent):
-        return ascent.sample_like(self.alpha, self.distribution).copysign_(ascent)
-
-class ParamSign(OptimizerModule):
-    """uses params with ascent sign"""
-    def __init__(self):
-        super().__init__({})
-
-
-    def _update(self, vars, ascent):
-        params = self.get_params()
-
-        return params.copysign(ascent)
-
-class NegParamSign(OptimizerModule):
-    """uses max(params_abs) - params_abs with ascent sign"""
-    def __init__(self):
-        super().__init__({})
-
-
-    def _update(self, vars, ascent):
-        neg_params = self.get_params().abs()
-        max = neg_params.total_max()
-        neg_params = neg_params.neg_().add(max)
-        return neg_params.copysign_(ascent)
-
-class InvParamSign(OptimizerModule):
-    """uses 1/(params_abs+eps) with ascent sign"""
-    def __init__(self, eps=1e-2):
-        super().__init__({})
-        self.eps = eps
-
-
-    def _update(self, vars, ascent):
-        inv_params = self.get_params().abs().add_(self.eps).reciprocal_()
-        return inv_params.copysign(ascent)
-
-
-class ParamWhereConsistentSign(OptimizerModule):
-    """where ascent and param signs are the same, it sets ascent to param value"""
-    def __init__(self, eps=1e-2):
-        super().__init__({})
-        self.eps = eps
-
-
-    def _update(self, vars, ascent):
-        params = self.get_params()
-        same_sign = params.sign() == ascent.sign()
-        ascent.masked_set_(same_sign, params)
-
-        return ascent

torchzero/modules/experimental/quad_interp.py

@@ -1,104 +0,0 @@
-import bisect
-
-import numpy as np
-import torch
-
-from ...tensorlist import TensorList
-from ...core import OptimizationVars
-from ..line_search.base_ls import LineSearchBase
-
-_FloatOrTensor = float | torch.Tensor
-
-def _ensure_float(x):
-    if isinstance(x, torch.Tensor): return x.detach().cpu().item()
-    elif isinstance(x, np.ndarray): return x.item()
-    return float(x)
-
-class Point:
-    def __init__(self, x, fx, dfx = None):
-        self.x = x
-        self.fx = fx
-        self.dfx = dfx
-
-    def __repr__(self):
-        return f'Point(x={self.x:.2f}, fx={self.fx:.2f})'
-
-def _step_2poins(x1, f1, df1, x2, f2):
-    # we have two points and one derivative
-    # minimize the quadratic to obtain 3rd point and perform bracketing
-    a = (df1 * x2 - f2 - df1*x1 + f1) / (x1**2 - x2**2 - 2*x1**2 + 2*x1*x2)
-    b = df1 - 2*a*x1
-    # c = -(a*x1**2 + b*x1 - y1)
-    return -b / (2 * a), a
-
-class QuadraticInterpolation2Point(LineSearchBase):
-    """This is WIP, please don't use yet!
-    Use `torchzero.modules.MinimizeQuadraticLS` and `torchzero.modules.MinimizeQuadratic3PointsLS` instead.
-
-    Args:
-        lr (_type_, optional): _description_. Defaults to 1e-2.
-        log_lrs (bool, optional): _description_. Defaults to False.
-        max_evals (int, optional): _description_. Defaults to 2.
-        min_dist (_type_, optional): _description_. Defaults to 1e-2.
-    """
-    def __init__(self, lr=1e-2, log_lrs = False, max_evals = 2, min_dist = 1e-2,):
-        super().__init__({"lr": lr}, maxiter=None, log_lrs=log_lrs)
-        self.max_evals = max_evals
-        self.min_dist = min_dist
-
-    @torch.no_grad
-    def _find_best_lr(self, vars: OptimizationVars, params: TensorList) -> float:
-        if vars.closure is None: raise ValueError('QuardaticLS requires closure')
-        closure = vars.closure
-        if vars.fx0 is None: vars.fx0 = vars.closure(False)
-        grad = vars.grad
-        if grad is None: grad = vars.ascent # in case we used FDM
-        if grad is None: raise ValueError('QuardaticLS requires gradients.')
-
-        params = self.get_params()
-        lr: float = self.get_first_group_key('lr') # this doesn't support variable lrs but we still want to support schedulers
-
-        # directional f'(x0)
-        # for each lr we step by this much
-        dfx0 = magn = grad.total_vector_norm(2)
-
-        # f(x1)
-        fx1 = self._evaluate_lr_(lr, closure, grad, params)
-
-        # make min_dist relative
-        min_dist = abs(lr) * self.min_dist
-        points = sorted([Point(0, _ensure_float(vars.fx0), dfx0), Point(lr, _ensure_float(fx1))], key = lambda x: x.fx)
-
-        for i in range(self.max_evals):
-            # find new point
-            p1, p2 = points
-            if p1.dfx is None: p1, p2 = p2, p1
-            xmin, curvature = _step_2poins(p1.x * magn, p1.fx, -p1.dfx, p2.x * magn, p2.fx) # type:ignore
-            xmin = _ensure_float(xmin/magn)
-            print(f'{xmin = }', f'{curvature = }, n_evals = {i+1}')
-
-            # if max_evals = 1, we just minimize a quadratic once
-            if i == self.max_evals - 1:
-                if curvature > 0: return xmin
-                return lr
-
-            # TODO: handle negative curvature
-            # if curvature < 0:
-            #     if points[0].x == 0: return lr
-            #     return points[0].x
-
-            # evaluate value and gradients at new point
-            fxmin = self._evaluate_lr_(xmin, closure, grad, params, backward=True)
-            dfxmin = -(params.grad * grad).total_sum()
-
-            # insort new point
-            bisect.insort(points, Point(xmin, _ensure_float(fxmin), dfxmin), key = lambda x: x.fx)
-
-            # pick 2 best points to find the new bracketing interval
-            points = sorted(points, key = lambda x: x.fx)[:2]
-            # TODO: new point might be worse than 2 existing ones which would lead to stagnation
-
-            # if points are too close, end the loop
-            if abs(points[0].x - points[1].x) < min_dist: break
-
-        return points[0].x

torchzero/modules/experimental/subspace.py

@@ -1,259 +0,0 @@
-import typing as T
-from abc import ABC, abstractmethod
-from collections import abc
-
-import torch
-
-from ... import tensorlist as tl
-from ...core import OptimizationVars, OptimizerModule, _Chain, _maybe_pass_backward
-# this whole thing can also be implemented via parameter vectors.
-# Need to test which one is more efficient...
-
-class Projection(ABC):
-    n = 1
-    @abstractmethod
-    def sample(self, params: tl.TensorList, vars: OptimizationVars) -> list[tl.TensorList]:
-        """Generate a projection.
-
-        Args:
-            params (tl.TensorList): tensor list of parameters.
-            state (OptimizationState): optimization state object.
-
-        Returns:
-            projection.
-        """
-
-class ProjRandom(Projection):
-    def __init__(self, n = 1, distribution: tl.Distributions = 'normal', ):
-        self.distribution: tl.Distributions = distribution
-        self.n = n
-
-    def sample(self, params: tl.TensorList, vars: OptimizationVars):
-        return [params.sample_like(distribution=self.distribution) for _ in range(self.n)]
-
-
-class Proj2Masks(Projection):
-    def __init__(self, n_pairs = 1):
-        """Similar to ProjRandom, but generates pairs of two random masks of 0s and 1s,
-        where second mask is an inverse of the first mask."""
-        self.n_pairs = n_pairs
-
-    @property
-    def n(self):
-        return self.n_pairs * 2
-
-    def sample(self, params: tl.TensorList, vars: OptimizationVars):
-        projections = []
-        for i in range(self.n_pairs):
-            mask = params.bernoulli_like(0.5)
-            mask2 = 1 - mask
-            projections.append(mask)
-            projections.append(mask2)
-
-        return projections
-
-
-class ProjAscent(Projection):
-    """Use ascent direction as the projection."""
-    def sample(self, params: tl.TensorList, vars: OptimizationVars):
-        if vars.ascent is None: raise ValueError
-        return [vars.ascent]
-
-class ProjAscentRay(Projection):
-    def __init__(self, eps = 0.1, n = 1, distribution: tl.Distributions = 'normal', ):
-        self.eps = eps
-        self.distribution: tl.Distributions = distribution
-        self.n = n
-
-    def sample(self, params: tl.TensorList, vars: OptimizationVars):
-        if vars.ascent is None: raise ValueError
-        mean = params.total_mean().detach().cpu().item()
-        return [vars.ascent + vars.ascent.sample_like(mean * self.eps, distribution=self.distribution) for _ in range(self.n)]
-
-class ProjGrad(Projection):
-    def sample(self, params: tl.TensorList, vars: OptimizationVars):
-        grad = vars.maybe_compute_grad_(params)
-        return [grad]
-
-class ProjGradRay(Projection):
-    def __init__(self, eps = 0.1, n = 1, distribution: tl.Distributions = 'normal', ):
-        self.eps = eps
-        self.distribution: tl.Distributions = distribution
-        self.n = n
-
-    def sample(self, params: tl.TensorList, vars: OptimizationVars):
-        grad = vars.maybe_compute_grad_(params)
-        mean = params.total_mean().detach().cpu().item()
-        return [grad + grad.sample_like(mean * self.eps, distribution=self.distribution) for _ in range(self.n)]
-
-class ProjGradAscentDifference(Projection):
-    def __init__(self, normalize=False):
-        """Use difference between gradient and ascent direction as projection.
-
-        Args:
-            normalize (bool, optional): normalizes grads and ascent projection to have norm = 1. Defaults to False.
-        """
-        self.normalize = normalize
-
-    def sample(self, params: tl.TensorList, vars: OptimizationVars):
-        grad = vars.maybe_compute_grad_(params)
-        if self.normalize:
-            return [vars.ascent / vars.ascent.total_vector_norm(2) - grad / grad.total_vector_norm(2)] # type:ignore
-
-        return [vars.ascent - grad] # type:ignore
-
-class ProjLastGradDifference(Projection):
-    def __init__(self):
-        """Use difference between last two gradients as the projection."""
-        self.last_grad = None
-    def sample(self, params: tl.TensorList, vars: OptimizationVars):
-        if self.last_grad is None:
-            self.last_grad = vars.maybe_compute_grad_(params)
-            return [self.last_grad]
-
-        grad = vars.maybe_compute_grad_(params)
-        diff = grad - self.last_grad
-        self.last_grad = grad
-        return [diff]
-
-class ProjLastAscentDifference(Projection):
-    def __init__(self):
-        """Use difference between last two ascent directions as the projection."""
-        self.last_direction = T.cast(tl.TensorList, None)
-
-    def sample(self, params: tl.TensorList, vars: OptimizationVars):
-        if self.last_direction is None:
-            self.last_direction: tl.TensorList = vars.ascent # type:ignore
-            return [self.last_direction]
-
-        diff = vars.ascent - self.last_direction # type:ignore
-        self.last_direction = vars.ascent # type:ignore
-        return [diff]
-
-class ProjNormalize(Projection):
-    def __init__(self, *projections: Projection):
-        """Normalizes all projections to have norm = 1."""
-        self.projections = projections
-
-    @property
-    def n(self):
-        return sum(proj.n for proj in self.projections)
-
-    def sample(self, params: tl.TensorList, vars: OptimizationVars): # type:ignore
-        vecs = [proj for obj in self.projections for proj in obj.sample(params, vars)]
-        norms = [v.total_vector_norm(2) for v in vecs]
-        return [v/norm if norm!=0 else v.randn_like() for v,norm in zip(vecs,norms)] # type:ignore
-
-class Subspace(OptimizerModule):
-    """This is pretty inefficient, I thought of a much better way to do this via jvp and I will rewrite this soon.
-
-    Optimizes parameters projected into a lower (or higher) dimensional subspace.
-
-    The subspace is a bunch of projections that go through the current point. Projections can be random,
-    or face in the direction of the gradient, or difference between last two gradients, etc. The projections
-    are updated every `update_every` steps.
-
-    Notes:
-        This doesn't work with anything that directly calculates the hessian or other quantities via `torch.autograd.grad`,
-        like `ExactNewton`. I will have to manually implement a subspace version for it.
-
-        This also zeroes parameters after each step, meaning it won't work with some integrations like nevergrad
-        (as they store their own parameters which don't get zeroed). It does however work with integrations like
-        `scipy.optimize` because they performs a full minimization on each step.
-    Another version of this which doesn't zero the params is under way.
-
-    Args:
-        projections (Projection | Iterable[Projection]):
-            list of projections - `Projection` objects that define the directions of the projections.
-            Each Projection object may generate one or multiple directions.
-        update_every (int, optional): generates new projections every n steps. Defaults to 1.
-    """
-    def __init__(
-        self,
-        modules: OptimizerModule | abc.Iterable[OptimizerModule],
-        projections: Projection | abc.Iterable[Projection],
-        update_every: int | None = 1,
-    ):
-        super().__init__({})
-        if isinstance(projections, Projection): projections = [projections]
-        self.projections = list(projections)
-        self._set_child_('subspace', modules)
-        self.update_every = update_every
-        self.current_step = 0
-
-        # cast them because they are guaranteed to be assigned on 1st step.
-        self.projection_vectors = T.cast(list[tl.TensorList], None)
-        self.projected_params = T.cast(torch.Tensor, None)
-
-
-    def _update_child_params_(self, child: "OptimizerModule"):
-        dtype = self._params[0].dtype
-        device = self._params[0].device
-        params = [torch.zeros(sum(proj.n for proj in self.projections), dtype = dtype, device = device, requires_grad=True)]
-        if child._has_custom_params: raise RuntimeError(f"Subspace child {child.__class__.__name__} can't have custom params.")
-        if not child._initialized:
-            child._initialize_(params, set_passed_params=False)
-        else:
-            child.param_groups = []
-            child.add_param_group({"params": params})
-
-    @torch.no_grad
-    def step(self, vars):
-        #if self.next_module is None: raise ValueError('RandomProjection needs a child')
-        if vars.closure is None: raise ValueError('RandomProjection needs a closure')
-        closure = vars.closure
-        params = self.get_params()
-
-        # every `regenerate_every` steps we generate new random projections.
-        if self.current_step == 0 or (self.update_every is not None and self.current_step % self.update_every == 0):
-
-            # generate n projection vetors
-            self.projection_vectors = [sample for proj in self.projections for sample in proj.sample(params, vars)]
-
-            # child params is n scalars corresponding to each projection vector
-            self.projected_params = self.children['subspace']._params[0] # type:ignore
-
-        # closure that takes the projected params from the child, puts them into full space params, and evaluates the loss
-        def projected_closure(backward = True):
-            residual = sum(vec * p for vec, p in zip(self.projection_vectors, self.projected_params))
-
-            # this in-place operation prevents autodiff from working
-            # we manually calculate the gradients as they are just a product
-            # therefore we need torch.no_grad here because optimizers call closure under torch.enabled_grad
-            with torch.no_grad(): params.add_(residual)
-
-            loss = _maybe_pass_backward(closure, backward)
-
-            if backward:
-                self.projected_params.grad = torch.cat([(params.grad * vec).total_sum().unsqueeze(0) for vec in self.projection_vectors])
-            with torch.no_grad(): params.sub_(residual)
-            return loss
-
-        # # if ascent direction is provided,
-        # # project the ascent direction into the projection space (need to test if this works)
-        # if ascent_direction is not None:
-        #     ascent_direction = tl.sum([ascent_direction*v for v in self.projection_vectors])
-
-        # perform a step with the child
-        subspace_state = vars.copy(False)
-        subspace_state.closure = projected_closure
-        subspace_state.ascent = None
-        if subspace_state.grad is not None:
-            subspace_state.grad = tl.TensorList([torch.cat([(params.grad * vec).total_sum().unsqueeze(0) for vec in self.projection_vectors])])
-        self.children['subspace'].step(subspace_state) # type:ignore
-
-        # that is going to update child's paramers, which we now project back to the full parameter space
-        residual = tl.sum([vec * p for vec, p in zip(self.projection_vectors, self.projected_params)])
-        vars.ascent = residual.neg_()
-
-        # move fx0 and fx0 approx to state
-        if subspace_state.fx0 is not None: vars.fx0 = subspace_state.fx0
-        if subspace_state.fx0_approx is not None: vars.fx0 = subspace_state.fx0_approx
-        # projected_params are residuals that have been applied to actual params on previous step in some way
-        # therefore they need to now become zero (otherwise they work like momentum with no decay).
-        # note: THIS WON'T WORK WITH INTEGRATIONS, UNLESS THEY PERFORM FULL MINIMIZATION EACH STEP
-        # because their params won't be zeroed.
-        self.projected_params.zero_()
-
-        self.current_step += 1
-        return self._update_params_or_step_with_next(vars)

torchzero/modules/gradient_approximation/__init__.py

@@ -1,7 +0,0 @@
-r"""
-Gradient approximation methods.
-"""
-from .fdm import FDM
-from .rfdm import RandomizedFDM
-from .newton_fdm import NewtonFDM
-from .forward_gradient import ForwardGradient

torchzero/modules/gradient_approximation/_fd_formulas.py

@@ -1,3 +0,0 @@
-import typing as T
-
-_FD_Formulas = T.Literal['central', 'forward', 'backward']