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

torchzero/modules/second_order/newton.py

@@ -44,8 +44,8 @@ def eigh_solve(H: torch.Tensor, g: torch.Tensor, tfm: Callable | None, search_ne
              # use eigvec or -eigvec depending on if it points in same direction as gradient
             return g.dot(d).sign() * d
 
-        L.reciprocal_()
-        return torch.linalg.multi_dot([Q * L.unsqueeze(-2), Q.mH, g]) # pylint:disable=not-callable
+        return Q @ ((Q.mH @ g) / L)
+
     except torch.linalg.LinAlgError:
         return None
 
@@ -53,46 +53,109 @@ def tikhonov_(H: torch.Tensor, reg: float):
     if reg!=0: H.add_(torch.eye(H.size(-1), dtype=H.dtype, device=H.device).mul_(reg))
     return H
 
-def eig_tikhonov_(H: torch.Tensor, reg: float):
-    v = torch.linalg.eigvalsh(H).min().clamp_(max=0).neg_() + reg # pylint:disable=not-callable
-    return tikhonov_(H, v)
-
 
 class Newton(Module):
-    """Exact newton via autograd.
+    """Exact newton's method via autograd.
+
+    .. 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.
+
 
     Args:
         reg (float, optional): tikhonov regularizer value. Defaults to 1e-6.
-        eig_reg (bool, optional): whether to use largest negative eigenvalue as regularizer. Defaults to False.
         search_negative (bool, Optional):
-            if True, whenever a negative eigenvalue is detected, the direction is taken along an eigenvector corresponding to a negative eigenvalue.
+            if True, whenever a negative eigenvalue is detected,
+            search direction is proposed along an eigenvector corresponding to a negative eigenvalue.
         hessian_method (str):
             how to calculate hessian. Defaults to "autograd".
         vectorize (bool, optional):
             whether to enable vectorized hessian. Defaults to True.
-        inner (Chainable | None, optional): inner modules. Defaults to None.
+        inner (Chainable | None, optional): modules to apply hessian preconditioner to. Defaults to None.
         H_tfm (Callable | None, optional):
             optional hessian transforms, takes in two arguments - `(hessian, gradient)`.
 
-            must return a tuple: `(hessian, is_inverted)` with transformed hessian and a boolean value
-            which must be True if transform inverted the hessian and False otherwise. Defaults to None.
+            must return either a tuple: `(hessian, is_inverted)` with transformed hessian and a boolean value
+            which must be True if transform inverted the hessian and False otherwise.
+
+            Or it returns a single tensor which is used as the update.
+
+            Defaults to None.
         eigval_tfm (Callable | None, optional):
             optional eigenvalues transform, for example :code:`torch.abs` or :code:`lambda L: torch.clip(L, min=1e-8)`.
-            If this is specified, eigendecomposition will be used to solve Hx = g.
+            If this is specified, eigendecomposition will be used to invert the hessian.
+
+    Examples:
+        Newton's method with backtracking line search
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Newton(),
+                tz.m.Backtracking()
+            )
+
+        Newton's method modified for non-convex functions by taking matrix absolute value of the hessian
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Newton(eigval_tfm=lambda x: torch.abs(x).clip(min=0.1)),
+                tz.m.Backtracking()
+            )
+
+        Newton's method modified for non-convex functions by searching along negative curvature directions
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Newton(search_negative=True),
+                tz.m.Backtracking()
+            )
+
+        Newton preconditioning applied to momentum
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Newton(inner=tz.m.EMA(0.9)),
+                tz.m.LR(0.1)
+            )
+
+        Diagonal newton example. This will still evaluate the entire hessian so it isn't efficient, but if you wanted to see how diagonal newton behaves or compares to full newton, you can use this.
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Newton(H_tfm = lambda H, g: g/H.diag()),
+                tz.m.Backtracking()
+            )
 
     """
     def __init__(
         self,
         reg: float = 1e-6,
-        eig_reg: bool = False,
         search_negative: bool = False,
+        update_freq: int = 1,
         hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
         vectorize: bool = True,
         inner: Chainable | None = None,
-        H_tfm: Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, bool]] | None = None,
+        H_tfm: Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, bool]] | Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
         eigval_tfm: Callable[[torch.Tensor], torch.Tensor] | None = None,
     ):
-        defaults = dict(reg=reg, eig_reg=eig_reg, hessian_method=hessian_method, vectorize=vectorize, H_tfm=H_tfm, eigval_tfm=eigval_tfm, search_negative=search_negative)
+        defaults = dict(reg=reg, hessian_method=hessian_method, vectorize=vectorize, H_tfm=H_tfm, eigval_tfm=eigval_tfm, search_negative=search_negative, update_freq=update_freq)
         super().__init__(defaults)
 
         if inner is not None:
@@ -106,47 +169,66 @@ class Newton(Module):
 
         settings = self.settings[params[0]]
         reg = settings['reg']
-        eig_reg = settings['eig_reg']
         search_negative = settings['search_negative']
         hessian_method = settings['hessian_method']
         vectorize = settings['vectorize']
         H_tfm = settings['H_tfm']
         eigval_tfm = settings['eigval_tfm']
+        update_freq = settings['update_freq']
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        g_list = var.grad
+        H = None
+        if step % update_freq == 0:
+            # ------------------------ calculate grad and hessian ------------------------ #
+            if hessian_method == 'autograd':
+                with torch.enable_grad():
+                    loss = var.loss = var.loss_approx = closure(False)
+                    g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
+                    g_list = [t[0] for t in g_list] # remove leading dim from loss
+                    var.grad = g_list
+                    H = hessian_list_to_mat(H_list)
 
-        # ------------------------ calculate grad and hessian ------------------------ #
-        if hessian_method == 'autograd':
-            with torch.enable_grad():
-                loss = var.loss = var.loss_approx = closure(False)
-                g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
-                g_list = [t[0] for t in g_list] # remove leading dim from loss
-                var.grad = g_list
-                H = hessian_list_to_mat(H_list)
-
-        elif hessian_method in ('func', 'autograd.functional'):
-            strat = 'forward-mode' if vectorize else 'reverse-mode'
-            with torch.enable_grad():
-                g_list = var.get_grad(retain_graph=True)
-                H: torch.Tensor = hessian_mat(partial(closure, backward=False), params,
-                                method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
-
-        else:
-            raise ValueError(hessian_method)
+            elif hessian_method in ('func', 'autograd.functional'):
+                strat = 'forward-mode' if vectorize else 'reverse-mode'
+                with torch.enable_grad():
+                    g_list = var.get_grad(retain_graph=True)
+                    H = hessian_mat(partial(closure, backward=False), params,
+                                    method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
+
+            else:
+                raise ValueError(hessian_method)
+
+            H = tikhonov_(H, reg)
+            if update_freq != 1:
+                self.global_state['H'] = H
+
+        if H is None:
+            H = self.global_state["H"]
+
+        # var.storage['hessian'] = H
 
         # -------------------------------- inner step -------------------------------- #
         update = var.get_update()
         if 'inner' in self.children:
-            update = apply_transform(self.children['inner'], update, params=params, grads=list(g_list), var=var)
+            update = apply_transform(self.children['inner'], update, params=params, grads=g_list, var=var)
+
         g = torch.cat([t.ravel() for t in update])
 
-        # ------------------------------- regulazition ------------------------------- #
-        if eig_reg: H = eig_tikhonov_(H, reg)
-        else: H = tikhonov_(H, reg)
 
         # ----------------------------------- solve ---------------------------------- #
         update = None
         if H_tfm is not None:
-            H, is_inv = H_tfm(H, g)
-            if is_inv: update = H @ g
+            ret = H_tfm(H, g)
+
+            if isinstance(ret, torch.Tensor):
+                update = ret
+
+            else: # returns (H, is_inv)
+                H, is_inv = ret
+                if is_inv: update = H @ g
 
         if search_negative or (eigval_tfm is not None):
             update = eigh_solve(H, g, eigval_tfm, search_negative=search_negative)
@@ -156,4 +238,101 @@ class Newton(Module):
         if update is None: update = least_squares_solve(H, g)
 
         var.update = vec_to_tensors(update, params)
+
+        return var
+
+class InverseFreeNewton(Module):
+    """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.
+    """
+    def __init__(
+        self,
+        update_freq: int = 1,
+        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
+        vectorize: bool = True,
+        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)
+
+    @torch.no_grad
+    def step(self, var):
+        params = TensorList(var.params)
+        closure = var.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        hessian_method = settings['hessian_method']
+        vectorize = settings['vectorize']
+        update_freq = settings['update_freq']
+
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        g_list = var.grad
+        Y = None
+        if step % update_freq == 0:
+            # ------------------------ calculate grad and hessian ------------------------ #
+            if hessian_method == 'autograd':
+                with torch.enable_grad():
+                    loss = var.loss = var.loss_approx = closure(False)
+                    g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
+                    g_list = [t[0] for t in g_list] # remove leading dim from loss
+                    var.grad = g_list
+                    H = hessian_list_to_mat(H_list)
+
+            elif hessian_method in ('func', 'autograd.functional'):
+                strat = 'forward-mode' if vectorize else 'reverse-mode'
+                with torch.enable_grad():
+                    g_list = var.get_grad(retain_graph=True)
+                    H = hessian_mat(partial(closure, backward=False), params,
+                                    method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
+
+            else:
+                raise ValueError(hessian_method)
+
+            # 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
+                eps = torch.finfo(H.dtype).eps
+                Y = self.global_state['Y'] = num.div_(denom.clip(min=eps, max=1/eps))
+
+            else:
+                Y = self.global_state['Y']
+                I = torch.eye(Y.size(0), device=Y.device, dtype=Y.dtype).mul_(2)
+                I -= H @ Y
+                Y = self.global_state['Y'] = Y @ I
+
+        if Y is None:
+            Y = self.global_state["Y"]
+
+
+        # -------------------------------- inner step -------------------------------- #
+        update = var.get_update()
+        if 'inner' in self.children:
+            update = apply_transform(self.children['inner'], update, params=params, grads=g_list, var=var)
+
+        g = torch.cat([t.ravel() for t in update])
+
+
+        # ----------------------------------- solve ---------------------------------- #
+        var.update = vec_to_tensors(Y@g, params)
+
         return var

torchzero/modules/second_order/newton_cg.py

@@ -1,26 +1,102 @@
-from collections.abc import Callable
 from typing import Literal, overload
-import warnings
 import torch
 
-from ...utils import TensorList, as_tensorlist, generic_zeros_like, generic_vector_norm, generic_numel
+from ...utils import TensorList, as_tensorlist, NumberList
 from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
 
 from ...core import Chainable, apply_transform, Module
-from ...utils.linalg.solve import cg
+from ...utils.linalg.solve import cg, steihaug_toint_cg, minres
 
 class NewtonCG(Module):
+    """Newton's method with a matrix-free conjugate gradient or minimial-residual solver.
+
+    This optimizer implements Newton's method using a matrix-free conjugate
+    gradient (CG) or a minimal-residual (MINRES) solver to approximate the search direction. Instead of
+    forming the full Hessian matrix, it only requires Hessian-vector products
+    (HVPs). These can be calculated efficiently using automatic
+    differentiation or approximated using finite differences.
+
+    .. note::
+        In most cases NewtonCG 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 HVPs.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    .. warning::
+        CG may fail if hessian is not positive-definite.
+
+    Args:
+        maxiter (int | None, optional):
+            Maximum number of iterations for the conjugate gradient solver.
+            By default, this is set to the number of dimensions in the
+            objective function, which is the theoretical upper bound for CG
+            convergence. Setting this to a smaller value (truncated Newton)
+            can still generate good search directions. Defaults to None.
+        tol (float, optional):
+            Relative tolerance for the conjugate gradient solver to determine
+            convergence. Defaults to 1e-4.
+        reg (float, optional):
+            Regularization parameter (damping) added to the Hessian diagonal.
+            This helps ensure the system is positive-definite. Defaults to 1e-8.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        h (float, optional):
+            The step size for finite differences if :code:`hvp_method` is
+            ``"forward"`` or ``"central"``. Defaults to 1e-3.
+        warm_start (bool, optional):
+            If ``True``, the conjugate gradient solver is initialized with the
+            solution from the previous optimization step. This can accelerate
+            convergence, especially in truncated Newton methods.
+            Defaults to False.
+        inner (Chainable | None, optional):
+            NewtonCG will attempt to apply preconditioning to the output of this module.
+
+    Examples:
+        Newton-CG with a backtracking line search:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.NewtonCG(),
+                tz.m.Backtracking()
+            )
+
+        Truncated Newton method (useful for large-scale problems):
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.NewtonCG(maxiter=10, warm_start=True),
+                tz.m.Backtracking()
+            )
+
+
+    """
     def __init__(
         self,
-        maxiter=None,
-        tol=1e-4,
+        maxiter: int | None = None,
+        tol: float = 1e-4,
         reg: float = 1e-8,
-        hvp_method: Literal["forward", "central", "autograd"] = "forward",
-        h=1e-3,
+        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        solver: Literal['cg', 'minres', 'minres_npc'] = 'cg',
+        h: float = 1e-3,
         warm_start=False,
         inner: Chainable | None = None,
     ):
-        defaults = dict(tol=tol, maxiter=maxiter, reg=reg, hvp_method=hvp_method, h=h, warm_start=warm_start)
+        defaults = dict(tol=tol, maxiter=maxiter, reg=reg, hvp_method=hvp_method, solver=solver, h=h, warm_start=warm_start)
         super().__init__(defaults,)
 
         if inner is not None:
@@ -37,6 +113,7 @@ class NewtonCG(Module):
         reg = settings['reg']
         maxiter = settings['maxiter']
         hvp_method = settings['hvp_method']
+        solver = settings['solver'].lower().strip()
         h = settings['h']
         warm_start = settings['warm_start']
 
@@ -68,13 +145,25 @@ class NewtonCG(Module):
         # -------------------------------- inner step -------------------------------- #
         b = var.get_update()
         if 'inner' in self.children:
-            b = as_tensorlist(apply_transform(self.children['inner'], b, params=params, grads=grad, var=var))
+            b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
+        b = as_tensorlist(b)
 
         # ---------------------------------- run cg ---------------------------------- #
         x0 = None
         if warm_start: x0 = self.get_state(params, 'prev_x', cls=TensorList) # initialized to 0 which is default anyway
 
-        x = cg(A_mm=H_mm, b=as_tensorlist(b), x0_=x0, tol=tol, maxiter=maxiter, reg=reg)
+        if solver == 'cg':
+            x = cg(A_mm=H_mm, b=b, x0_=x0, tol=tol, maxiter=maxiter, reg=reg)
+
+        elif solver == 'minres':
+            x = minres(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=False)
+
+        elif solver == 'minres_npc':
+            x = minres(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=True)
+
+        else:
+            raise ValueError(f"Unknown solver {solver}")
+
         if warm_start:
             assert x0 is not None
             x0.copy_(x)
@@ -83,3 +172,203 @@ class NewtonCG(Module):
         return var
 
 
+class TruncatedNewtonCG(Module):
+    """Trust region Newton's method with a matrix-free Steihaug-Toint conjugate gradient or MINRES solver.
+
+    This optimizer implements Newton's method using a matrix-free conjugate
+    gradient (CG) solver to approximate the search direction. Instead of
+    forming the full Hessian matrix, it only requires Hessian-vector products
+    (HVPs). These can be calculated efficiently using automatic
+    differentiation or approximated using finite differences.
+
+    .. note::
+        In most cases NewtonCGSteihaug 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 HVPs.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    .. warning::
+        CG may fail if hessian is not positive-definite.
+
+    Args:
+        maxiter (int | None, optional):
+            Maximum number of iterations for the conjugate gradient solver.
+            By default, this is set to the number of dimensions in the
+            objective function, which is the theoretical upper bound for CG
+            convergence. Setting this to a smaller value (truncated Newton)
+            can still generate good search directions. Defaults to None.
+        eta (float, optional):
+            whenever actual to predicted loss reduction ratio is larger than this, a step is accepted.
+        nplus (float, optional):
+            trust region multiplier on successful steps.
+        nminus (float, optional):
+            trust region multiplier on unsuccessful steps.
+        init (float, optional): initial trust region.
+        tol (float, optional):
+            Relative tolerance for the conjugate gradient solver to determine
+            convergence. Defaults to 1e-4.
+        reg (float, optional):
+            Regularization parameter (damping) added to the Hessian diagonal.
+            This helps ensure the system is positive-definite. Defaults to 1e-8.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        h (float, optional):
+            The step size for finite differences if :code:`hvp_method` is
+            ``"forward"`` or ``"central"``. Defaults to 1e-3.
+        inner (Chainable | None, optional):
+            NewtonCG will attempt to apply preconditioning to the output of this module.
+
+    Examples:
+        Trust-region Newton-CG:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.NewtonCGSteihaug(),
+            )
+
+    Reference:
+        Steihaug, Trond. "The conjugate gradient method and trust regions in large scale optimization." SIAM Journal on Numerical Analysis 20.3 (1983): 626-637.
+    """
+    def __init__(
+        self,
+        maxiter: int | None = None,
+        eta: float= 1e-6,
+        nplus: float = 2,
+        nminus: float = 0.25,
+        init: float = 1,
+        tol: float = 1e-4,
+        reg: float = 1e-8,
+        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        solver: Literal['cg', 'minres', 'minres_npc'] = 'cg',
+        h: float = 1e-3,
+        max_attempts: int = 10,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(tol=tol, maxiter=maxiter, reg=reg, hvp_method=hvp_method, h=h, eta=eta, nplus=nplus, nminus=nminus, init=init, max_attempts=max_attempts, solver=solver)
+        super().__init__(defaults,)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, var):
+        params = TensorList(var.params)
+        closure = var.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        tol = settings['tol']
+        reg = settings['reg']
+        maxiter = settings['maxiter']
+        hvp_method = settings['hvp_method']
+        h = settings['h']
+        max_attempts = settings['max_attempts']
+        solver = settings['solver'].lower().strip()
+
+        eta = settings['eta']
+        nplus = settings['nplus']
+        nminus = settings['nminus']
+        init = settings['init']
+
+        # ---------------------- Hessian vector product function --------------------- #
+        if hvp_method == 'autograd':
+            grad = var.get_grad(create_graph=True)
+
+            def H_mm(x):
+                with torch.enable_grad():
+                    return TensorList(hvp(params, grad, x, retain_graph=True))
+
+        else:
+
+            with torch.enable_grad():
+                grad = var.get_grad()
+
+            if hvp_method == 'forward':
+                def H_mm(x):
+                    return TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1])
+
+            elif hvp_method == 'central':
+                def H_mm(x):
+                    return TensorList(hvp_fd_central(closure, params, x, h=h, normalize=True)[1])
+
+            else:
+                raise ValueError(hvp_method)
+
+
+        # -------------------------------- inner step -------------------------------- #
+        b = var.get_update()
+        if 'inner' in self.children:
+            b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
+        b = as_tensorlist(b)
+
+        # ---------------------------------- run cg ---------------------------------- #
+        success = False
+        x = None
+        while not success:
+            max_attempts -= 1
+            if max_attempts < 0: break
+
+            trust_region = self.global_state.get('trust_region', init)
+            if trust_region < 1e-8 or trust_region > 1e8:
+                trust_region = self.global_state['trust_region'] = init
+
+            if solver == 'cg':
+                x = steihaug_toint_cg(A_mm=H_mm, b=b, trust_region=trust_region, tol=tol, maxiter=maxiter, reg=reg)
+
+            elif solver == 'minres':
+                x = minres(A_mm=H_mm, b=b, trust_region=trust_region, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=False)
+
+            elif solver == 'minres_npc':
+                x = minres(A_mm=H_mm, b=b, trust_region=trust_region, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=True)
+
+            else:
+                raise ValueError(f"unknown solver {solver}")
+
+            # ------------------------------- trust region ------------------------------- #
+            Hx = H_mm(x)
+            pred_reduction = b.dot(x) - 0.5 * x.dot(Hx)
+
+            params -= x
+            loss_star = closure(False)
+            params += x
+            reduction = var.get_loss(False) - loss_star
+
+            rho = reduction / (pred_reduction.clip(min=1e-8))
+
+            # failed step
+            if rho < 0.25:
+                self.global_state['trust_region'] = trust_region * nminus
+
+            # very good step
+            elif rho > 0.75:
+                diff = trust_region - x.abs()
+                if (diff.global_min() / trust_region) > 1e-4: # hits boundary
+                    self.global_state['trust_region'] = trust_region * nplus
+
+            # if the ratio is high enough then accept the proposed step
+            if rho > eta:
+                success = True
+
+        assert x is not None
+        if success:
+            var.update = x
+
+        else:
+            var.update = params.zeros_like()
+
+        return var
+
+