torchzero 0.3.15__py3-none-any.whl → 0.4.1__py3-none-any.whl

torchzero/modules/weight_decay/reinit.py

@@ -0,0 +1,83 @@
+from functools import partial
+
+import torch
+
+from ...core import Module
+from ...utils import NumberList, TensorList
+
+
+def _reset_except_self(objective, modules, self: Module):
+    for m in modules:
+        if m is not self:
+            m.reset()
+
+class RandomReinitialize(Module):
+    """On each step with probability ``p_reinit`` trigger reinitialization,
+    whereby ``p_weights`` weights are reset to their initial values.
+
+    This modifies the parameters directly. Place it as the first module.
+
+    Args:
+        p_reinit (float, optional): probability to trigger reinitialization on each step. Defaults to 0.01.
+        p_weights (float, optional): probability for each weight to be set to initial value when reinitialization is triggered. Defaults to 0.1.
+        store_every (int | None, optional): if set, stores new initial values every this many steps. Defaults to None.
+        beta (float, optional):
+            whenever ``store_every`` is triggered, uses linear interpolation with this beta.
+            If ``store_every=1``, this can be set to some value close to 1 such as 0.999
+            to reinitialize to slow parameter EMA. Defaults to 0.
+        reset (bool, optional): whether to reset states of other modules on reinitialization. Defaults to False.
+        seed (int | None, optional): random seed.
+    """
+
+    def __init__(
+        self,
+        p_reinit: float = 0.01,
+        p_weights: float = 0.1,
+        store_every: int | None = None,
+        beta: float = 0,
+        reset: bool = False,
+        seed: int | None = None,
+    ):
+        defaults = dict(p_weights=p_weights, p_reinit=p_reinit, store_every=store_every, beta=beta, reset=reset, seed=seed)
+        super().__init__(defaults)
+
+    def update(self, objective):
+        # this stores initial values to per-parameter states
+        p_init = self.get_state(objective.params, "p_init", init="params", cls=TensorList)
+
+        # store new params every store_every steps
+        step = self.global_state.get("step", 0)
+        self.global_state["step"] = step + 1
+
+        store_every = self.defaults["store_every"]
+        if (store_every is not None and step % store_every == 0):
+            beta = self.get_settings(objective.params, "beta", cls=NumberList)
+            p_init.lerp_(objective.params, weight=(1 - beta))
+
+    @torch.no_grad
+    def apply(self, objective):
+        p_reinit = self.defaults["p_reinit"]
+        device = objective.params[0].device
+        generator = self.get_generator(device, self.defaults["seed"])
+
+        # determine whether to trigger reinitialization
+        reinitialize = torch.rand(1, generator=generator, device=device) < p_reinit
+
+        # reinitialize
+        if reinitialize:
+            params = TensorList(objective.params)
+            p_init = self.get_state(params, "p_init", init=params)
+
+
+            # mask with p_weights entries being True
+            p_weights = self.get_settings(params, "p_weights")
+            mask = params.bernoulli_like(p_weights, generator=generator).as_bool()
+
+            # set weights at mask to their initialization
+            params.masked_set_(mask, p_init)
+
+            # reset
+            if self.defaults["reset"]:
+                objective.post_step_hooks.append(partial(_reset_except_self, self=self))
+
+        return objective

torchzero/modules/weight_decay/weight_decay.py

@@ -3,7 +3,7 @@ from typing import Literal
 
 import torch
 
-from ...core import Module, Target, Transform
+from ...core import Module,  TensorTransform
 from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states, Metrics
 
 
@@ -21,7 +21,7 @@ def weight_decay_(
     return grad_.add_(params.pow(ord-1).copysign_(params).mul_(weight_decay))
 
 
-class WeightDecay(Transform):
+class WeightDecay(TensorTransform):
     """Weight decay.
 
     Args:
@@ -33,7 +33,7 @@ class WeightDecay(Transform):
 
     Adam with non-decoupled weight decay
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.WeightDecay(1e-3),
         tz.m.Adam(),
@@ -44,7 +44,7 @@ class WeightDecay(Transform):
     Adam with decoupled weight decay that still scales with learning rate
     ```python
 
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adam(),
         tz.m.WeightDecay(1e-3),
@@ -54,7 +54,7 @@ class WeightDecay(Transform):
 
     Adam with fully decoupled weight decay that doesn't scale with learning rate
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adam(),
         tz.m.LR(1e-3),
@@ -63,19 +63,19 @@ class WeightDecay(Transform):
     ```
 
     """
-    def __init__(self, weight_decay: float, ord: int = 2, target: Target = 'update'):
+    def __init__(self, weight_decay: float, ord: int = 2):
 
         defaults = dict(weight_decay=weight_decay, ord=ord)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         weight_decay = NumberList(s['weight_decay'] for s in settings)
         ord = settings[0]['ord']
 
         return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay, ord)
 
-class RelativeWeightDecay(Transform):
+class RelativeWeightDecay(TensorTransform):
     """Weight decay relative to the mean absolute value of update, gradient or parameters depending on value of ``norm_input`` argument.
 
     Args:
@@ -93,7 +93,7 @@ class RelativeWeightDecay(Transform):
 
     Adam with non-decoupled relative weight decay
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.RelativeWeightDecay(1e-1),
         tz.m.Adam(),
@@ -103,7 +103,7 @@ class RelativeWeightDecay(Transform):
 
     Adam with decoupled relative weight decay
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Adam(),
         tz.m.RelativeWeightDecay(1e-1),
@@ -117,13 +117,12 @@ class RelativeWeightDecay(Transform):
         ord: int  = 2,
         norm_input: Literal["update", "grad", "params"] = "update",
         metric: Metrics = 'mad',
-        target: Target = "update",
     ):
         defaults = dict(weight_decay=weight_decay, ord=ord, norm_input=norm_input, metric=metric)
-        super().__init__(defaults, uses_grad=norm_input == 'grad', target=target)
+        super().__init__(defaults, uses_grad=norm_input == 'grad')
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         weight_decay = NumberList(s['weight_decay'] for s in settings)
 
         ord = settings[0]['ord']
@@ -161,9 +160,9 @@ class DirectWeightDecay(Module):
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, var):
-        weight_decay = self.get_settings(var.params, 'weight_decay', cls=NumberList)
+    def apply(self, objective):
+        weight_decay = self.get_settings(objective.params, 'weight_decay', cls=NumberList)
         ord = self.defaults['ord']
 
-        decay_weights_(var.params, weight_decay, ord)
-        return var
+        decay_weights_(objective.params, weight_decay, ord)
+        return objective

torchzero/modules/wrappers/optim_wrapper.py

@@ -3,7 +3,7 @@ from typing import Any
 import torch
 
 from ...core.module import Module
-from ...utils import Params, _copy_param_groups, _make_param_groups
+from ...utils.params import Params, _copy_param_groups, _make_param_groups
 
 
 class Wrap(Module):
@@ -11,7 +11,7 @@ class Wrap(Module):
     Wraps a pytorch optimizer to use it as a module.
 
     Note:
-        Custom param groups are supported only by ``set_param_groups``, settings passed to Modular will be applied to all parameters.
+        Custom param groups are supported only by ``set_param_groups``, settings passed to Optimizer will be applied to all parameters.
 
     Args:
         opt_fn (Callable[..., torch.optim.Optimizer] | torch.optim.Optimizer):
@@ -21,7 +21,7 @@ class Wrap(Module):
         **kwargs:
             Extra args to be passed to opt_fn. The function is called as ``opt_fn(parameters, *args, **kwargs)``.
         use_param_groups:
-            Whether to pass settings passed to Modular to the wrapped optimizer.
+            Whether to pass settings passed to Optimizer to the wrapped optimizer.
 
             Note that settings to the first parameter are used for all parameters,
             so if you specified per-parameter settings, they will be ignored.
@@ -32,7 +32,7 @@ class Wrap(Module):
     ```python
 
     from pytorch_optimizer import StableAdamW
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Wrap(StableAdamW, lr=1),
         tz.m.Cautious(),
@@ -66,8 +66,8 @@ class Wrap(Module):
         return super().set_param_groups(param_groups)
 
     @torch.no_grad
-    def step(self, var):
-        params = var.params
+    def apply(self, objective):
+        params = objective.params
 
         # initialize opt on 1st step
         if self.optimizer is None:
@@ -76,14 +76,14 @@ class Wrap(Module):
             self.optimizer = self._opt_fn(param_groups, *self._opt_args, **self._opt_kwargs)
 
         # set optimizer per-parameter settings
-        if self.defaults["use_param_groups"] and var.modular is not None:
+        if self.defaults["use_param_groups"] and objective.modular is not None:
             for group in self.optimizer.param_groups:
                 first_param = group['params'][0]
                 setting = self.settings[first_param]
 
                 # settings passed in `set_param_groups` are the highest priority
                 # schedulers will override defaults but not settings passed in `set_param_groups`
-                # this is consistent with how Modular does it.
+                # this is consistent with how Optimizer does it.
                 if self._custom_param_groups is not None:
                     setting = {k:v for k,v in setting if k not in self._custom_param_groups[0]}
 
@@ -91,19 +91,19 @@ class Wrap(Module):
 
         # set grad to update
         orig_grad = [p.grad for p in params]
-        for p, u in zip(params, var.get_update()):
+        for p, u in zip(params, objective.get_updates()):
             p.grad = u
 
         # if this is last module, simply use optimizer to update parameters
-        if var.modular is not None and self is var.modular.modules[-1]:
+        if objective.modular is not None and self is objective.modular.modules[-1]:
             self.optimizer.step()
 
             # restore grad
             for p, g in zip(params, orig_grad):
                 p.grad = g
 
-            var.stop = True; var.skip_update = True
-            return var
+            objective.stop = True; objective.skip_update = True
+            return objective
 
         # this is not the last module, meaning update is difference in parameters
         # and passed to next module
@@ -111,11 +111,11 @@ class Wrap(Module):
         self.optimizer.step() # step and update params
         for p, g in zip(params, orig_grad):
             p.grad = g
-        var.update = list(torch._foreach_sub(params_before_step, params)) # set update to difference between params
+        objective.updates = list(torch._foreach_sub(params_before_step, params)) # set update to difference between params
         for p, o in zip(params, params_before_step):
             p.set_(o) # pyright: ignore[reportArgumentType]
 
-        return var
+        return objective
 
     def reset(self):
         super().reset()

torchzero/modules/zeroth_order/cd.py

@@ -29,17 +29,20 @@ class CD(Module):
             whether to use three points (three function evaluatins) to determine descent direction.
             if False, uses two points, but then ``adaptive`` can't be used. Defaults to True.
     """
-    def __init__(self, h:float=1e-3, grad:bool=True, adaptive:bool=True, index:Literal['cyclic', 'cyclic2', 'random']="cyclic2", threepoint:bool=True,):
+    def __init__(self, h:float=1e-3, grad:bool=False, adaptive:bool=True, index:Literal['cyclic', 'cyclic2', 'random']="cyclic2", threepoint:bool=True,):
         defaults = dict(h=h, grad=grad, adaptive=adaptive, index=index, threepoint=threepoint)
         super().__init__(defaults)
 
+    def update(self, objective): raise RuntimeError
+    def apply(self, objective): raise RuntimeError
+
     @torch.no_grad
-    def step(self, var):
-        closure = var.closure
+    def step(self, objective):
+        closure = objective.closure
         if closure is None:
             raise RuntimeError("CD requires closure")
 
-        params = TensorList(var.params)
+        params = TensorList(objective.params)
         ndim = params.global_numel()
 
         grad_step_size = self.defaults['grad']
@@ -79,7 +82,7 @@ class CD(Module):
             else:
                 warnings.warn("CD adaptive=True only works with threepoint=True")
 
-        f_0 = var.get_loss(False)
+        f_0 = objective.get_loss(False)
         params.flat_set_lambda_(idx, lambda x: x + h)
         f_p = closure(False)
 
@@ -117,6 +120,6 @@ class CD(Module):
         # ----------------------------- create the update ---------------------------- #
         update = params.zeros_like()
         update.flat_set_(idx, alpha)
-        var.update = update
-        return var
+        objective.updates = update
+        return objective
 

torchzero/optim/mbs.py

@@ -0,0 +1,291 @@
+from typing import NamedTuple
+import math
+from collections.abc import Iterable
+from decimal import ROUND_HALF_UP, Decimal
+
+import numpy as np
+
+
+def format_number(number, n):
+    """Rounds to n significant digits after the decimal point."""
+    if number == 0: return 0
+    if math.isnan(number) or math.isinf(number) or (not math.isfinite(number)): return number
+    if n <= 0: raise ValueError("n must be positive")
+
+    dec = Decimal(str(number))
+    if dec.is_zero(): return 0
+    if number > 10**n or dec % 1 == 0: return int(dec)
+
+    if abs(dec) >= 1:
+        places = n
+    else:
+        frac_str = format(abs(dec), 'f').split('.')[1]
+        leading_zeros = len(frac_str) - len(frac_str.lstrip('0'))
+        places = leading_zeros + n
+
+    quantizer = Decimal('1e-' + str(places))
+    rounded_dec = dec.quantize(quantizer, rounding=ROUND_HALF_UP)
+
+    if rounded_dec % 1 == 0: return int(rounded_dec)
+    return float(rounded_dec)
+
+def _nonfinite_to_inf(x):
+    if not math.isfinite(x): return math.inf
+    return x
+
+def _tofloatlist(x) -> list[float]:
+    if isinstance(x, (int,float)): return [x]
+    if isinstance(x, np.ndarray) and x.size == 1: return [float(x.item())]
+    return [float(i) for i in x]
+
+class Trial(NamedTuple):
+    x: float
+    f: tuple[float, ...]
+
+class Solution(NamedTuple):
+    x: float
+    f: tuple[float, ...]
+    trials: list[Trial]
+
+class MBS:
+    """Univariate minimization via grid search followed by refining, supports multi-objective functions.
+
+    This tends to outperform bayesian optimization for learning rate tuning, it is also good for plotting.
+
+    First it evaluates all points defined in ``grid``. The grid doesn't have to be dense and the solution doesn't
+    have to be between the endpoints.
+
+    Then it picks ``num_candidates`` best points per each objective. If any of those points are endpoints,
+    it expands the search space by ``step`` in that direction and evaluates the new endpoint.
+
+    Otherwise it keeps picking points between best points and evaluating them, until ``num_binary`` evaluations
+    have been performed.
+
+    Args:
+        grid (Iterable[float], optional): values for initial grid search. If ``log_scale=True``, should be in log10 scale.
+        step (float, optional): expansion step size. Defaults to 1.
+        num_candidates (int, optional): number of best points to sample new points around on each iteration. Defaults to 2.
+        num_binary (int, optional): maximum number of new points sampled via binary search. Defaults to 7.
+        num_expansions (int, optional): maximum number of expansions (not counted towards binary search points). Defaults to 7.
+        rounding (int, optional): rounding is to significant digits, avoids evaluating points that are too close.
+        lb (float | None, optional): lower bound. If ``log_scale=True``, should be in log10 scale.
+        ub (float | None, optional): upper bound. If ``log_scale=True``, should be in log10 scale.
+        log_scale (bool, optional):
+            whether to minimize in log10 scale. If true, it is assumed that
+            ``grid``, ``lb`` and ``ub`` are given in log10 scale.
+
+    Example:
+
+    ```python
+    def objective(x: float):
+        x = x * 4
+        return -(np.sin(x) * (x / 3) + np.cos(x*2.5) * 2 - 0.05 * (x-5)**2)
+
+    mbs = MBS(grid=[-1, 0, 1, 2, 3, 4], step=1, num_binary=10, num_expansions=10)
+
+    x, f, trials = mbs.run(objective)
+    # x - solution
+    # f - value at solution x
+    # trials - list of trials, each trial is a named tuple: Trial(x, f)
+    """
+
+    def __init__(
+        self,
+        grid: Iterable[float],
+        step: float,
+        num_candidates: int = 3,
+        num_binary: int = 20,
+        num_expansions: int = 20,
+        rounding: int| None = 2,
+        lb = None,
+        ub = None,
+        log_scale: bool = False,
+    ):
+        self.objectives: dict[int, dict[float,float]] = {}
+        """dictionary of objectives, each maps point (x) to value (v)"""
+
+        self.evaluated: set[float] = set()
+        """set of evaluated points (x)"""
+
+        grid = tuple(grid)
+        if len(grid) == 0: raise ValueError("At least one grid search point must be specified")
+        self.grid = sorted(grid)
+
+        self.step = step
+        self.num_candidates = num_candidates
+        self.num_binary = num_binary
+        self.num_expansions = num_expansions
+        self.rounding = rounding
+        self.log_scale = log_scale
+        self.lb = lb
+        self.ub = ub
+
+    def _get_best_x(self, n: int, objective: int):
+        """n best points"""
+        obj = self.objectives[objective]
+        v_to_x = [(v,x) for x,v in obj.items()]
+        v_to_x.sort(key = lambda vx: vx[0])
+        xs = [x for v,x in v_to_x]
+        return xs[:n]
+
+    def _suggest_points_around(self, x: float, objective: int):
+        """suggests points around x"""
+        points = list(self.objectives[objective].keys())
+        points.sort()
+        if x not in points: raise RuntimeError(f"{x} not in {points}")
+
+        expansions = []
+        if x == points[0]:
+            expansions.append((x-self.step, 'expansion'))
+
+        if x == points[-1]:
+            expansions.append((x+self.step, 'expansion'))
+
+        if len(expansions) != 0: return expansions
+
+        idx = points.index(x)
+        xm = points[idx-1]
+        xp = points[idx+1]
+
+        x1 = (x - (x - xm)/2)
+        x2 = (x + (xp - x)/2)
+
+        return [(x1, 'binary'), (x2, 'binary')]
+
+    def _out_of_bounds(self, x):
+        if self.lb is not None and x < self.lb: return True
+        if self.ub is not None and x > self.ub: return True
+        return False
+
+    def _evaluate(self, fn, x):
+        """Evaluate a point, returns False if point is already in history"""
+        if self.rounding is not None: x = format_number(x, self.rounding)
+        if x in self.evaluated: return False
+        if self._out_of_bounds(x): return False
+
+        self.evaluated.add(x)
+
+        if self.log_scale: vals = _tofloatlist(fn(10 ** x))
+        else: vals = _tofloatlist(fn(x))
+        vals = [_nonfinite_to_inf(v) for v in vals]
+
+        for idx, v in enumerate(vals):
+            if idx not in self.objectives: self.objectives[idx] = {}
+            self.objectives[idx][x] = v
+
+        return True
+
+    def run(self, fn) -> Solution:
+        # step 1 - gr id search
+        for x in self.grid:
+            self._evaluate(fn, x)
+
+        # step 2 - binary search
+        while True:
+            if (self.num_candidates <= 0) or (self.num_expansions <= 0 and self.num_binary <= 0): break
+
+            # suggest candidates
+            candidates: list[tuple[float, str]] = []
+
+            # sample around best points
+            for objective in self.objectives:
+                best_points = self._get_best_x(self.num_candidates, objective)
+                for p in best_points:
+                    candidates.extend(self._suggest_points_around(p, objective=objective))
+
+            # filter
+            if self.num_expansions <= 0:
+                candidates = [(x,t) for x,t in candidates if t != 'expansion']
+
+            if self.num_candidates <= 0:
+                candidates = [(x,t) for x,t in candidates if t != 'binary']
+
+            # if expansion was suggested, discard anything else
+            types = [t for x, t in candidates]
+            if any(t == 'expansion' for t in types):
+                candidates = [(x,t) for x,t in candidates if t == 'expansion']
+
+            # evaluate candidates
+            terminate = False
+            at_least_one_evaluated = False
+            for x, t in candidates:
+                evaluated = self._evaluate(fn, x)
+                if not evaluated: continue
+                at_least_one_evaluated = True
+
+                if t == 'expansion': self.num_expansions -= 1
+                elif t == 'binary': self.num_binary -= 1
+
+                if self.num_binary < 0:
+                    terminate = True
+                    break
+
+            if terminate: break
+            if not at_least_one_evaluated:
+                if self.rounding is None: break
+                self.rounding += 1
+                if self.rounding == 100: break
+
+        # create dict[float, tuple[float,...]]
+        ret = {}
+        for i, objective in enumerate(self.objectives.values()):
+            for x, v in objective.items():
+                if self.log_scale: x = 10 ** x
+                if x not in ret: ret[x] = [None for _ in self.objectives]
+                ret[x][i] = v
+
+        for v in ret.values():
+            assert len(v) == len(self.objectives), v
+            assert all(i is not None for i in v), v
+
+        # ret maps x to list of per-objective values, e.g. {1: [0.1, 0.3], ...}
+        # now make a list of trials as they are easier to work with
+        trials: list[Trial] = []
+        for x, values in ret.items():
+            trials.append(Trial(x=x, f=values))
+
+        # sort trials by sum of values
+        trials.sort(key = lambda trial: sum(trial.f))
+        return Solution(x=trials[0].x, f=trials[0].f, trials=trials)
+
+def mbs_minimize(
+    fn,
+    grid: Iterable[float],
+    step: float,
+    num_candidates: int = 3,
+    num_binary: int = 20,
+    num_expansions: int = 20,
+    rounding=2,
+    lb:float | None = None,
+    ub:float | None = None,
+    log_scale=False,
+) -> Solution:
+    """minimize univariate function via MBS.
+
+    Args:
+        fn (function): objective function that accepts a float and returns a float or a sequence of floats to minimize.
+        step (float, optional): expansion step size. Defaults to 1.
+        num_candidates (int, optional): number of best points to sample new points around on each iteration. Defaults to 2.
+        num_binary (int, optional): maximum number of new points sampled via binary search. Defaults to 7.
+        num_expansions (int, optional): maximum number of expansions (not counted towards binary search points). Defaults to 7.
+        rounding (int, optional): rounding is to significant digits, avoids evaluating points that are too close.
+        lb (float | None, optional): lower bound. If ``log_scale=True``, should be in log10 scale.
+        ub (float | None, optional): upper bound. If ``log_scale=True``, should be in log10 scale.
+        log_scale (bool, optional):
+            whether to minimize in log10 scale. If true, it is assumed that
+            ``grid``, ``lb`` and ``ub`` are given in log10 scale.
+
+    Example:
+
+    ```python
+    def objective(x: float):
+        x = x * 4
+        return -(np.sin(x) * (x / 3) + np.cos(x*2.5) * 2 - 0.05 * (x-5)**2)
+
+    x, f, trials = mbs_minimize(objective, grid=[-1, 0, 1, 2, 3, 4], step=1, num_binary=10, num_expansions=10)
+    # x - solution
+    # f - value at solution x
+    # trials - list of trials, each trial is a named tuple: Trial(x, f)
+    """
+    mbs = MBS(grid, step=step, num_candidates=num_candidates, num_binary=num_binary, num_expansions=num_expansions, rounding=rounding, lb=lb, ub=ub, log_scale=log_scale)
+    return mbs.run(fn)

torchzero/optim/root.py

@@ -3,7 +3,7 @@ from collections.abc import Callable
 
 from abc import abstractmethod
 import torch
-from ..modules.higher_order.multipoint import sixth_order_im1, sixth_order_p6, _solve
+from ..modules.second_order.multipoint import sixth_order_3p, sixth_order_5p, two_point_newton, sixth_order_3pm2, _solve
 
 def make_evaluate(f: Callable[[torch.Tensor], torch.Tensor]):
     def evaluate(x, order) -> tuple[torch.Tensor, ...]:
@@ -53,7 +53,7 @@ class Newton(RootBase):
     def one_iteration(self, x, evaluate): return newton(x, evaluate, self.lstsq)
 
 
-class SixthOrderP6(RootBase):
+class SixthOrder3P(RootBase):
     """sixth-order iterative method
 
     Abro, Hameer Akhtar, and Muhammad Mujtaba Shaikh. "A new time-efficient and convergent nonlinear solver." Applied Mathematics and Computation 355 (2019): 516-536.
@@ -62,4 +62,4 @@ class SixthOrderP6(RootBase):
     def one_iteration(self, x, evaluate):
         def f(x): return evaluate(x, 0)[0]
         def f_j(x): return evaluate(x, 1)
-        return sixth_order_p6(x, f, f_j, self.lstsq)
+        return sixth_order_3p(x, f, f_j, self.lstsq)

torchzero/optim/utility/split.py

@@ -3,7 +3,8 @@ from collections.abc import Callable, Iterable
 
 import torch
 
-from ...utils import flatten, get_params
+from ...utils import flatten
+from ...utils.optimizer import get_params
 
 class Split(torch.optim.Optimizer):
     """Steps will all `optimizers`, also has a check that they have no duplicate parameters.