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

torchzero/modules/restarts/restars.py

@@ -4,12 +4,14 @@ from typing import final, Literal, cast
 
 import torch
 
-from ...core import Chainable, Module, Var
+from ...core import Chainable, Module, Objective
 from ...utils import TensorList
 from ..termination import TerminationCriteriaBase
 
-def _reset_except_self(optimizer, var, self: Module):
-    for m in optimizer.unrolled_modules: m.reset()
+def _reset_except_self(objective, modules, self: Module):
+    for m in modules:
+        if m is not self:
+            m.reset()
 
 class RestartStrategyBase(Module, ABC):
     """Base class for restart strategies.
@@ -24,38 +26,38 @@ class RestartStrategyBase(Module, ABC):
             self.set_child('modules', modules)
 
     @abstractmethod
-    def should_reset(self, var: Var) -> bool:
+    def should_reset(self, objective: Objective) -> bool:
         """returns whether reset should occur"""
 
-    def _reset_on_condition(self, var):
+    def _reset_on_condition(self, objective: Objective):
         modules = self.children.get('modules', None)
 
-        if self.should_reset(var):
+        if self.should_reset(objective):
             if modules is None:
-                var.post_step_hooks.append(partial(_reset_except_self, self=self))
+                objective.post_step_hooks.append(partial(_reset_except_self, self=self))
             else:
                 modules.reset()
 
         return modules
 
     @final
-    def update(self, var):
-        modules = self._reset_on_condition(var)
+    def update(self, objective):
+        modules = self._reset_on_condition(objective)
         if modules is not None:
-            modules.update(var)
+            modules.update(objective)
 
     @final
-    def apply(self, var):
+    def apply(self, objective):
         # don't check here because it was check in `update`
         modules = self.children.get('modules', None)
-        if modules is None: return var
-        return modules.apply(var.clone(clone_update=False))
+        if modules is None: return objective
+        return modules.apply(objective.clone(clone_updates=False))
 
     @final
-    def step(self, var):
-        modules = self._reset_on_condition(var)
-        if modules is None: return var
-        return modules.step(var.clone(clone_update=False))
+    def step(self, objective):
+        modules = self._reset_on_condition(objective)
+        if modules is None: return objective
+        return modules.step(objective.clone(clone_updates=False))
 
 
 
@@ -76,11 +78,11 @@ class RestartOnStuck(RestartStrategyBase):
         super().__init__(defaults, modules)
 
     @torch.no_grad
-    def should_reset(self, var):
+    def should_reset(self, objective):
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
-        params = TensorList(var.params)
+        params = TensorList(objective.params)
         tol = self.defaults['tol']
         if tol is None: tol = torch.finfo(params[0].dtype).tiny * 2
         n_tol = self.defaults['n_tol']
@@ -122,12 +124,12 @@ class RestartEvery(RestartStrategyBase):
         defaults = dict(steps=steps)
         super().__init__(defaults, modules)
 
-    def should_reset(self, var):
+    def should_reset(self, objective):
         step = self.global_state.get('step', 0) + 1
         self.global_state['step'] = step
 
         n = self.defaults['steps']
-        if isinstance(n, str): n = sum(p.numel() for p in var.params if p.requires_grad)
+        if isinstance(n, str): n = sum(p.numel() for p in objective.params if p.requires_grad)
 
         # reset every n steps
         if step % n == 0:
@@ -141,9 +143,9 @@ class RestartOnTerminationCriteria(RestartStrategyBase):
         super().__init__(None, modules)
         self.set_child('criteria', criteria)
 
-    def should_reset(self, var):
+    def should_reset(self, objective):
         criteria = cast(TerminationCriteriaBase, self.children['criteria'])
-        return criteria.should_terminate(var)
+        return criteria.should_terminate(objective)
 
 class PowellRestart(RestartStrategyBase):
     """Powell's two restarting criterions for conjugate gradient methods.
@@ -169,14 +171,14 @@ class PowellRestart(RestartStrategyBase):
         defaults=dict(cond1=cond1, cond2=cond2)
         super().__init__(defaults, modules)
 
-    def should_reset(self, var):
-        g = TensorList(var.get_grad())
+    def should_reset(self, objective):
+        g = TensorList(objective.get_grads())
         cond1 = self.defaults['cond1']; cond2 = self.defaults['cond2']
 
         # -------------------------------- initialize -------------------------------- #
         if 'initialized' not in self.global_state:
             self.global_state['initialized'] = 0
-            g_prev = self.get_state(var.params, 'g_prev', init=g)
+            g_prev = self.get_state(objective.params, 'g_prev', init=g)
             return False
 
         g_g = g.dot(g)
@@ -184,7 +186,7 @@ class PowellRestart(RestartStrategyBase):
         reset = False
         # ------------------------------- 1st condition ------------------------------ #
         if cond1 is not None:
-            g_prev = self.get_state(var.params, 'g_prev', must_exist=True, cls=TensorList)
+            g_prev = self.get_state(objective.params, 'g_prev', must_exist=True, cls=TensorList)
             g_g_prev = g_prev.dot(g)
 
             if g_g_prev.abs() >= cond1 * g_g:
@@ -192,7 +194,7 @@ class PowellRestart(RestartStrategyBase):
 
         # ------------------------------- 2nd condition ------------------------------ #
         if (cond2 is not None) and (not reset):
-            d_g = TensorList(var.get_update()).dot(g)
+            d_g = TensorList(objective.get_updates()).dot(g)
             if (-1-cond2) * g_g < d_g < (-1 + cond2) * g_g:
                 reset = True
 
@@ -229,17 +231,17 @@ class BirginMartinezRestart(Module):
 
         self.set_child("module", module)
 
-    def update(self, var):
+    def update(self, objective):
         module = self.children['module']
-        module.update(var)
+        module.update(objective)
 
-    def apply(self, var):
+    def apply(self, objective):
         module = self.children['module']
-        var = module.apply(var.clone(clone_update=False))
+        objective = module.apply(objective.clone(clone_updates=False))
 
         cond = self.defaults['cond']
-        g = TensorList(var.get_grad())
-        d = TensorList(var.get_update())
+        g = TensorList(objective.get_grads())
+        d = TensorList(objective.get_updates())
         d_g = d.dot(g)
         d_norm = d.global_vector_norm()
         g_norm = g.global_vector_norm()
@@ -247,7 +249,7 @@ class BirginMartinezRestart(Module):
         # d in our case is same direction as g so it has a minus sign
         if -d_g > -cond * d_norm * g_norm:
             module.reset()
-            var.update = g.clone()
-            return var
+            objective.updates = g.clone()
+            return objective
 
-        return var
+        return objective

torchzero/modules/second_order/__init__.py

@@ -1,7 +1,7 @@
 from .ifn import InverseFreeNewton
-from .inm import INM
+from .inm import ImprovedNewton
 from .multipoint import SixthOrder3P, SixthOrder3PM2, SixthOrder5P, TwoPointNewton
 from .newton import Newton
 from .newton_cg import NewtonCG, NewtonCGSteihaug
 from .nystrom import NystromPCG, NystromSketchAndSolve
-from .rsn import RSN
+from .rsn import SubspaceNewton

torchzero/modules/second_order/ifn.py

@@ -1,89 +1,58 @@
-import warnings
-from collections.abc import Callable
-from functools import partial
-from typing import Literal
-
 import torch
 
-from ...core import Chainable, Module, apply_transform, Var
+from ...core import Chainable, Transform, HessianMethod
 from ...utils import TensorList, vec_to_tensors
-from ...utils.linalg.linear_operator import DenseWithInverse, Dense
-from .newton import _get_H, _get_loss_grad_and_hessian, _newton_step
+from ...linalg.linear_operator import DenseWithInverse
 
 
-class InverseFreeNewton(Module):
+class InverseFreeNewton(Transform):
     """Inverse-free newton's method
 
-    .. note::
-        In most cases Newton should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Newton preconditioning to another module's output.
-
-    .. note::
-        This module requires the a closure passed to the optimizer step,
-        as it needs to re-evaluate the loss and gradients for calculating the hessian.
-        The closure must accept a ``backward`` argument (refer to documentation).
-
-    .. warning::
-        this uses roughly O(N^2) memory.
-
     Reference
         [Massalski, Marcin, and Magdalena Nockowska-Rosiak. "INVERSE-FREE NEWTON'S METHOD." Journal of Applied Analysis & Computation 15.4 (2025): 2238-2257.](https://www.jaac-online.com/article/doi/10.11948/20240428)
     """
     def __init__(
         self,
         update_freq: int = 1,
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
+        hessian_method: HessianMethod = "batched_autograd",
+        h: float = 1e-3,
         inner: Chainable | None = None,
     ):
-        defaults = dict(hessian_method=hessian_method, vectorize=vectorize, update_freq=update_freq)
-        super().__init__(defaults)
-
-        if inner is not None:
-            self.set_child('inner', inner)
+        defaults = dict(hessian_method=hessian_method, h=h)
+        super().__init__(defaults, update_freq=update_freq, inner=inner)
 
     @torch.no_grad
-    def update(self, var):
-        update_freq = self.defaults['update_freq']
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
 
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
+        _, _, H = objective.hessian(
+            hessian_method=fs['hessian_method'],
+            h=fs['h'],
+            at_x0=True
+        )
 
-        if step % update_freq == 0:
-            loss, g_list, H = _get_loss_grad_and_hessian(
-                var, self.defaults['hessian_method'], self.defaults['vectorize']
-            )
-            self.global_state["H"] = H
+        self.global_state["H"] = H
 
-            # inverse free part
-            if 'Y' not in self.global_state:
-                num = H.T
-                denom = (torch.linalg.norm(H, 1) * torch.linalg.norm(H, float('inf'))) # pylint:disable=not-callable
+        # inverse free part
+        if 'Y' not in self.global_state:
+            num = H.T
+            denom = (torch.linalg.norm(H, 1) * torch.linalg.norm(H, float('inf'))) # pylint:disable=not-callable
 
-                finfo = torch.finfo(H.dtype)
-                self.global_state['Y'] = num.div_(denom.clip(min=finfo.tiny * 2, max=finfo.max / 2))
+            finfo = torch.finfo(H.dtype)
+            self.global_state['Y'] = num.div_(denom.clip(min=finfo.tiny * 2, max=finfo.max / 2))
 
-            else:
-                Y = self.global_state['Y']
-                I2 = torch.eye(Y.size(0), device=Y.device, dtype=Y.dtype).mul_(2)
-                I2 -= H @ Y
-                self.global_state['Y'] = Y @ I2
+        else:
+            Y = self.global_state['Y']
+            I2 = torch.eye(Y.size(0), device=Y.device, dtype=Y.dtype).mul_(2)
+            I2 -= H @ Y
+            self.global_state['Y'] = Y @ I2
 
 
-    def apply(self, var):
+    def apply_states(self, objective, states, settings):
         Y = self.global_state["Y"]
-        params = var.params
-
-        # -------------------------------- inner step -------------------------------- #
-        update = var.get_update()
-        if 'inner' in self.children:
-            update = apply_transform(self.children['inner'], update, params=params, grads=var.grad, var=var)
-
-        g = torch.cat([t.ravel() for t in update])
-
-        # ----------------------------------- solve ---------------------------------- #
-        var.update = vec_to_tensors(Y@g, params)
-
-        return var
+        g = torch.cat([t.ravel() for t in objective.get_updates()])
+        objective.updates = vec_to_tensors(Y@g, objective.params)
+        return objective
 
-    def get_H(self,var):
+    def get_H(self,objective=...):
         return DenseWithInverse(A = self.global_state["H"], A_inv=self.global_state["Y"])

torchzero/modules/second_order/inm.py

@@ -1,12 +1,11 @@
 from collections.abc import Callable
-from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module
-from ...utils import TensorList, vec_to_tensors
-from ..functional import safe_clip
-from .newton import _get_H, _get_loss_grad_and_hessian, _newton_step
+from ...core import Chainable, Transform, HessianMethod
+from ...utils import TensorList, vec_to_tensors_, unpack_states
+from ..opt_utils import safe_clip
+from .newton import _newton_update_state_, _newton_solve, _newton_get_H
 
 @torch.no_grad
 def inm(f:torch.Tensor, J:torch.Tensor, s:torch.Tensor, y:torch.Tensor):
@@ -25,7 +24,7 @@ def _eigval_fn(J: torch.Tensor, fn) -> torch.Tensor:
     L, Q = torch.linalg.eigh(J) # pylint:disable=not-callable
     return (Q * L.unsqueeze(-2)) @ Q.mH
 
-class INM(Module):
+class ImprovedNewton(Transform):
     """Improved Newton's Method (INM).
 
     Reference:
@@ -35,71 +34,76 @@ class INM(Module):
     def __init__(
         self,
         damping: float = 0,
-        use_lstsq: bool = False,
+        eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
         update_freq: int = 1,
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
+        precompute_inverse: bool | None = None,
+        use_lstsq: bool = False,
+        hessian_method: HessianMethod = "batched_autograd",
+        h: float = 1e-3,
         inner: Chainable | None = None,
-        H_tfm: Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, bool]] | Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
-        eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
     ):
-        defaults = dict(damping=damping, hessian_method=hessian_method, use_lstsq=use_lstsq, vectorize=vectorize, H_tfm=H_tfm, eigval_fn=eigval_fn, update_freq=update_freq)
-        super().__init__(defaults)
-
-        if inner is not None:
-            self.set_child("inner", inner)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner'], defaults["update_freq"]
+        super().__init__(defaults, update_freq=update_freq, inner=inner, )
 
     @torch.no_grad
-    def update(self, var):
-        update_freq = self.defaults['update_freq']
-
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
 
-        if step % update_freq == 0:
-            _, f_list, J = _get_loss_grad_and_hessian(
-                var, self.defaults['hessian_method'], self.defaults['vectorize']
-            )
+        _, f_list, J = objective.hessian(
+            hessian_method=fs['hessian_method'],
+            h=fs['h'],
+            at_x0=True
+        )
+        if f_list is None: f_list = objective.get_grads()
 
-            f = torch.cat([t.ravel() for t in f_list])
-            J = _eigval_fn(J, self.defaults["eigval_fn"])
+        f = torch.cat([t.ravel() for t in f_list])
+        J = _eigval_fn(J, fs["eigval_fn"])
 
-            x_list = TensorList(var.params)
-            f_list = TensorList(var.get_grad())
-            x_prev, f_prev = self.get_state(var.params, "x_prev", "f_prev", cls=TensorList)
+        x_list = TensorList(objective.params)
+        f_list = TensorList(objective.get_grads())
+        x_prev, f_prev = unpack_states(states, objective.params, "x_prev", "f_prev", cls=TensorList)
 
-            # initialize on 1st step, do Newton step
-            if step == 0:
-                x_prev.copy_(x_list)
-                f_prev.copy_(f_list)
-                self.global_state["P"] = J
-                return
+        # initialize on 1st step, do Newton step
+        if "H" not in self.global_state:
+            x_prev.copy_(x_list)
+            f_prev.copy_(f_list)
+            P = J
 
-            # INM update
+        # INM update
+        else:
             s_list = x_list - x_prev
             y_list = f_list - f_prev
             x_prev.copy_(x_list)
             f_prev.copy_(f_list)
 
-            self.global_state["P"] = inm(f, J, s=s_list.to_vec(), y=y_list.to_vec())
+            P = inm(f, J, s=s_list.to_vec(), y=y_list.to_vec())
 
+        # update state
+        precompute_inverse = fs["precompute_inverse"]
+        if precompute_inverse is None:
+            precompute_inverse = fs["__update_freq"] >= 10
 
-    @torch.no_grad
-    def apply(self, var):
-        params = var.params
-        update = _newton_step(
-            var=var,
-            H = self.global_state["P"],
-            damping=self.defaults["damping"],
-            inner=self.children.get("inner", None),
-            H_tfm=self.defaults["H_tfm"],
-            eigval_fn=None, # it is applied in `update`
-            use_lstsq=self.defaults["use_lstsq"],
+        _newton_update_state_(
+            H=P,
+            state = self.global_state,
+            damping = fs["damping"],
+            eigval_fn = fs["eigval_fn"],
+            precompute_inverse = precompute_inverse,
+            use_lstsq = fs["use_lstsq"]
         )
 
-        var.update = vec_to_tensors(update, params)
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        updates = objective.get_updates()
+        fs = settings[0]
+
+        b = torch.cat([t.ravel() for t in updates])
+        sol = _newton_solve(b=b, state=self.global_state, use_lstsq=fs["use_lstsq"])
+
+        vec_to_tensors_(sol, updates)
+        return objective
 
-        return var
 
-    def get_H(self,var=...):
-        return _get_H(self.global_state["P"], eigval_fn=None)
+    def get_H(self,objective=...):
+        return _newton_get_H(self.global_state)

torchzero/modules/second_order/multipoint.py

@@ -1,19 +1,17 @@
-from collections.abc import Callable
-from contextlib import nullcontext
 from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping
+from typing import Any
+
 import numpy as np
 import torch
 
-from ...core import Chainable, Module, apply_transform, Var
-from ...utils import TensorList, vec_to_tensors, vec_to_tensors_
-from ...utils.derivatives import (
-    flatten_jacobian,
-    jacobian_wrt,
-)
+from ...core import Chainable, DerivativesMethod, Objective, Transform
+from ...utils import TensorList, vec_to_tensors
+
 
-class HigherOrderMethodBase(Module, ABC):
-    def __init__(self, defaults: dict | None = None, vectorize: bool = True):
-        self._vectorize = vectorize
+class HigherOrderMethodBase(Transform, ABC):
+    def __init__(self, defaults: dict | None = None, derivatives_method: DerivativesMethod = 'batched_autograd'):
+        self._derivatives_method: DerivativesMethod = derivatives_method
         super().__init__(defaults)
 
     @abstractmethod
@@ -21,61 +19,27 @@ class HigherOrderMethodBase(Module, ABC):
         self,
         x: torch.Tensor,
         evaluate: Callable[[torch.Tensor, int], tuple[torch.Tensor, ...]],
-        var: Var,
+        objective: Objective,
+        setting: Mapping[str, Any],
     ) -> torch.Tensor:
         """"""
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        x0 = params.clone()
-        closure = var.closure
+    def apply_states(self, objective, states, settings):
+        params = TensorList(objective.params)
+
+        closure = objective.closure
         if closure is None: raise RuntimeError('MultipointNewton requires closure')
-        vectorize = self._vectorize
+        derivatives_method = self._derivatives_method
 
         def evaluate(x, order) -> tuple[torch.Tensor, ...]:
             """order=0 - returns (loss,), order=1 - returns (loss, grad), order=2 - returns (loss, grad, hessian), etc."""
-            params.from_vec_(x)
-
-            if order == 0:
-                loss = closure(False)
-                params.copy_(x0)
-                return (loss, )
-
-            if order == 1:
-                with torch.enable_grad():
-                    loss = closure()
-                grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
-                params.copy_(x0)
-                return loss, torch.cat([g.ravel() for g in grad])
-
-            with torch.enable_grad():
-                loss = var.loss = var.loss_approx = closure(False)
-
-                g_list = torch.autograd.grad(loss, params, create_graph=True)
-                var.grad = list(g_list)
-
-                g = torch.cat([t.ravel() for t in g_list])
-                n = g.numel()
-                ret = [loss, g]
-                T = g # current derivatives tensor
-
-                # get all derivative up to order
-                for o in range(2, order + 1):
-                    is_last = o == order
-                    T_list = jacobian_wrt([T], params, create_graph=not is_last, batched=vectorize)
-                    with torch.no_grad() if is_last else nullcontext():
-                        # the shape is (ndim, ) * order
-                        T = flatten_jacobian(T_list).view(n, n, *T.shape[1:])
-                        ret.append(T)
-
-            params.copy_(x0)
-            return tuple(ret)
+            return objective.derivatives_at(x, order, method=derivatives_method)
 
         x = torch.cat([p.ravel() for p in params])
-        dir = self.one_iteration(x, evaluate, var)
-        var.update = vec_to_tensors(dir, var.params)
-        return var
+        dir = self.one_iteration(x, evaluate, objective, settings[0])
+        objective.updates = vec_to_tensors(dir, objective.params)
+        return objective
 
 def _inv(A: torch.Tensor, lstsq:bool) -> torch.Tensor:
     if lstsq: return torch.linalg.pinv(A) # pylint:disable=not-callable
@@ -106,16 +70,15 @@ class SixthOrder3P(HigherOrderMethodBase):
 
     Abro, Hameer Akhtar, and Muhammad Mujtaba Shaikh. "A new time-efficient and convergent nonlinear solver." Applied Mathematics and Computation 355 (2019): 516-536.
     """
-    def __init__(self, lstsq: bool=False, vectorize: bool = True):
+    def __init__(self, lstsq: bool=False, derivatives_method: DerivativesMethod = 'batched_autograd'):
         defaults=dict(lstsq=lstsq)
-        super().__init__(defaults=defaults, vectorize=vectorize)
+        super().__init__(defaults=defaults, derivatives_method=derivatives_method)
 
-    def one_iteration(self, x, evaluate, var):
-        settings = self.defaults
-        lstsq = settings['lstsq']
+    @torch.no_grad
+    def one_iteration(self, x, evaluate, objective, setting):
         def f(x): return evaluate(x, 1)[1]
         def f_j(x): return evaluate(x, 2)[1:]
-        x_star = sixth_order_3p(x, f, f_j, lstsq)
+        x_star = sixth_order_3p(x, f, f_j, setting['lstsq'])
         return x - x_star
 
 # I don't think it works (I tested root finding with this and it goes all over the place)
@@ -173,15 +136,14 @@ def sixth_order_5p(x:torch.Tensor, f_j, lstsq:bool=False):
 
 class SixthOrder5P(HigherOrderMethodBase):
     """Argyros, Ioannis K., et al. "Extended convergence for two sixth order methods under the same weak conditions." Foundations 3.1 (2023): 127-139."""
-    def __init__(self, lstsq: bool=False, vectorize: bool = True):
+    def __init__(self, lstsq: bool=False, derivatives_method: DerivativesMethod = 'batched_autograd'):
         defaults=dict(lstsq=lstsq)
-        super().__init__(defaults=defaults, vectorize=vectorize)
+        super().__init__(defaults=defaults, derivatives_method=derivatives_method)
 
-    def one_iteration(self, x, evaluate, var):
-        settings = self.defaults
-        lstsq = settings['lstsq']
+    @torch.no_grad
+    def one_iteration(self, x, evaluate, objective, setting):
         def f_j(x): return evaluate(x, 2)[1:]
-        x_star = sixth_order_5p(x, f_j, lstsq)
+        x_star = sixth_order_5p(x, f_j, setting['lstsq'])
         return x - x_star
 
 # 2f 1J 2 solves
@@ -196,16 +158,15 @@ class TwoPointNewton(HigherOrderMethodBase):
     """two-point Newton method with frozen derivative with third order convergence.
 
     Sharma, Janak Raj, and Deepak Kumar. "A fast and efficient composite Newton–Chebyshev method for systems of nonlinear equations." Journal of Complexity 49 (2018): 56-73."""
-    def __init__(self, lstsq: bool=False, vectorize: bool = True):
+    def __init__(self, lstsq: bool=False, derivatives_method: DerivativesMethod = 'batched_autograd'):
         defaults=dict(lstsq=lstsq)
-        super().__init__(defaults=defaults, vectorize=vectorize)
+        super().__init__(defaults=defaults, derivatives_method=derivatives_method)
 
-    def one_iteration(self, x, evaluate, var):
-        settings = self.defaults
-        lstsq = settings['lstsq']
+    @torch.no_grad
+    def one_iteration(self, x, evaluate, objective, setting):
         def f(x): return evaluate(x, 1)[1]
         def f_j(x): return evaluate(x, 2)[1:]
-        x_star = two_point_newton(x, f, f_j, lstsq)
+        x_star = two_point_newton(x, f, f_j, setting['lstsq'])
         return x - x_star
 
 #3f 2J 1inv
@@ -224,15 +185,14 @@ def sixth_order_3pm2(x:torch.Tensor, f, f_j, lstsq:bool=False):
 
 class SixthOrder3PM2(HigherOrderMethodBase):
     """Wang, Xiaofeng, and Yang Li. "An efficient sixth-order Newton-type method for solving nonlinear systems." Algorithms 10.2 (2017): 45."""
-    def __init__(self, lstsq: bool=False, vectorize: bool = True):
+    def __init__(self, lstsq: bool=False, derivatives_method: DerivativesMethod = 'batched_autograd'):
         defaults=dict(lstsq=lstsq)
-        super().__init__(defaults=defaults, vectorize=vectorize)
+        super().__init__(defaults=defaults, derivatives_method=derivatives_method)
 
-    def one_iteration(self, x, evaluate, var):
-        settings = self.defaults
-        lstsq = settings['lstsq']
+    @torch.no_grad
+    def one_iteration(self, x, evaluate, objective, setting):
         def f_j(x): return evaluate(x, 2)[1:]
         def f(x): return evaluate(x, 1)[1]
-        x_star = sixth_order_3pm2(x, f, f_j, lstsq)
+        x_star = sixth_order_3pm2(x, f, f_j, setting['lstsq'])
         return x - x_star