torchzero 0.3.14__py3-none-any.whl → 0.4.0__py3-none-any.whl

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
 

torchzero/modules/second_order/newton.py

@@ -1,21 +1,12 @@
-import warnings
 from collections.abc import Callable
-from functools import partial
 from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, apply_transform
-from ...utils import TensorList, vec_to_tensors
-from ...utils.derivatives import (
-    flatten_jacobian,
-    hessian_mat,
-    hvp,
-    hvp_fd_central,
-    hvp_fd_forward,
-    jacobian_and_hessian_wrt,
-)
-from ...utils.linalg.linear_operator import DenseWithInverse, Dense
+from ...core import Chainable, Transform, Objective, HessianMethod, Module
+from ...utils import vec_to_tensors
+from ...linalg.linear_operator import Dense, DenseWithInverse
+
 
 def _lu_solve(H: torch.Tensor, g: torch.Tensor):
     try:
@@ -26,10 +17,9 @@ def _lu_solve(H: torch.Tensor, g: torch.Tensor):
         return None
 
 def _cholesky_solve(H: torch.Tensor, g: torch.Tensor):
-    x, info = torch.linalg.cholesky_ex(H) # pylint:disable=not-callable
+    L, info = torch.linalg.cholesky_ex(H) # pylint:disable=not-callable
     if info == 0:
-        g.unsqueeze_(1)
-        return torch.cholesky_solve(g, x)
+        return torch.cholesky_solve(g.unsqueeze(-1), L).squeeze(-1)
     return None
 
 def _least_squares_solve(H: torch.Tensor, g: torch.Tensor):
@@ -49,10 +39,59 @@ def _eigh_solve(H: torch.Tensor, g: torch.Tensor, tfm: Callable | None, search_n
     except torch.linalg.LinAlgError:
         return None
 
+def _newton_step(objective: Objective, H: torch.Tensor, damping:float, H_tfm, eigval_fn, use_lstsq:bool, g_proj: Callable | None = None, no_inner: Module | None = None) -> torch.Tensor:
+    """INNER SHOULD BE NONE IN MOST CASES! Because Transform already has inner.
+    Returns the update tensor, then do vec_to_tensor(update, params)"""
+    # -------------------------------- inner step -------------------------------- #
+    if no_inner is not None:
+        objective = no_inner.step(objective)
+
+    update = objective.get_updates()
+
+    g = torch.cat([t.ravel() for t in update])
+    if g_proj is not None: g = g_proj(g)
+
+    # ----------------------------------- solve ---------------------------------- #
+    update = None
+
+    if damping != 0:
+        H = H + torch.eye(H.size(-1), dtype=H.dtype, device=H.device).mul_(damping)
+
+    if H_tfm is not None:
+        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 eigval_fn is not None:
+        update = _eigh_solve(H, g, eigval_fn, search_negative=False)
+
+    if update is None and use_lstsq: update = _least_squares_solve(H, g)
+    if update is None: update = _cholesky_solve(H, g)
+    if update is None: update = _lu_solve(H, g)
+    if update is None: update = _least_squares_solve(H, g)
+
+    return update
+
+def _get_H(H: torch.Tensor, eigval_fn):
+    if eigval_fn is not None:
+        try:
+            L, Q = torch.linalg.eigh(H) # pylint:disable=not-callable
+            L: torch.Tensor = eigval_fn(L)
+            H = Q @ L.diag_embed() @ Q.mH
+            H_inv = Q @ L.reciprocal().diag_embed() @ Q.mH
+            return DenseWithInverse(H, H_inv)
 
+        except torch.linalg.LinAlgError:
+            pass
 
+    return Dense(H)
 
-class Newton(Module):
+class Newton(Transform):
     """Exact newton's method via autograd.
 
     Newton's method produces a direction jumping to the stationary point of quadratic approximation of the target function.
@@ -60,7 +99,7 @@ class Newton(Module):
     ``g`` can be output of another module, if it is specifed in ``inner`` argument.
 
     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.
+        In most cases Newton should be the first module in the chain because it relies on autograd. Use the ``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,
@@ -77,11 +116,6 @@ class Newton(Module):
             when hessian is not invertible. If False, tries cholesky, if it fails tries LU, and then least squares.
             If ``eigval_fn`` is specified, eigendecomposition will always be used to solve the linear system and this
             argument will be ignored.
-        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): modules to apply hessian preconditioner to. Defaults to None.
         H_tfm (Callable | None, optional):
             optional hessian transforms, takes in two arguments - `(hessian, gradient)`.
 
@@ -94,6 +128,22 @@ class Newton(Module):
         eigval_fn (Callable | None, optional):
             optional eigenvalues transform, for example ``torch.abs`` or ``lambda L: torch.clip(L, min=1e-8)``.
             If this is specified, eigendecomposition will be used to invert the hessian.
+        hessian_method (str):
+            Determines how hessian is computed.
+
+            - ``"batched_autograd"`` - uses autograd to compute ``ndim`` batched hessian-vector products. Faster than ``"autograd"`` but uses more memory.
+            - ``"autograd"`` - uses autograd to compute ``ndim`` hessian-vector products using for loop. Slower than ``"batched_autograd"`` but uses less memory.
+            - ``"functional_revrev"`` - uses ``torch.autograd.functional`` with "reverse-over-reverse" strategy and a for-loop. This is generally equivalent to ``"autograd"``.
+            - ``"functional_fwdrev"`` - uses ``torch.autograd.functional`` with vectorized "forward-over-reverse" strategy. Faster than ``"functional_fwdrev"`` but uses more memory (``"batched_autograd"`` seems to be faster)
+            - ``"func"`` - uses ``torch.func.hessian`` which uses "forward-over-reverse" strategy. This method is the fastest and is recommended, however it is more restrictive and fails with some operators which is why it isn't the default.
+            - ``"gfd_forward"`` - computes ``ndim`` hessian-vector products via gradient finite difference using a less accurate forward formula which requires one extra gradient evaluation per hessian-vector product.
+            - ``"gfd_central"`` - computes ``ndim`` hessian-vector products via gradient finite difference using a more accurate central formula which requires two gradient evaluations per hessian-vector product.
+            - ``"fd"`` - uses function values to estimate gradient and hessian via finite difference. This uses less evaluations than chaining ``"gfd_*"`` after ``tz.m.FDM``.
+
+            Defaults to ``"batched_autograd"``.
+        h (float, optional):
+            finite difference step size for "fd_forward" and "fd_central".
+        inner (Chainable | None, optional): modules to apply hessian preconditioner to. Defaults to None.
 
     # See also
 
@@ -111,10 +161,9 @@ class Newton(Module):
     The linear system is solved via cholesky decomposition, if that fails, LU decomposition, and if that fails, least squares.
     Least squares can be forced by setting ``use_lstsq=True``, which may generate better search directions when linear system is overdetermined.
 
-    Additionally, if ``eigval_fn`` is specified or ``search_negative`` is ``True``,
-    eigendecomposition of the hessian is computed, ``eigval_fn`` is applied to the eigenvalues,
-    and ``(H + yI)⁻¹`` is computed using the computed eigenvectors and transformed eigenvalues.
-    This is more generally more computationally expensive.
+    Additionally, if ``eigval_fn`` is specified, eigendecomposition of the hessian is computed,
+    ``eigval_fn`` is applied to the eigenvalues, and ``(H + yI)⁻¹`` is computed using the computed eigenvectors and transformed eigenvalues. This is more generally more computationally expensive,
+    but not by much
 
     ## Handling non-convexity
 
@@ -167,217 +216,45 @@ class Newton(Module):
     def __init__(
         self,
         damping: float = 0,
-        search_negative: bool = False,
         use_lstsq: 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]] | Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
         eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        hessian_method: HessianMethod = "batched_autograd",
+        h: float = 1e-3,
+        inner: Chainable | None = None,
     ):
-        defaults = dict(damping=damping, hessian_method=hessian_method, use_lstsq=use_lstsq, vectorize=vectorize, H_tfm=H_tfm, eigval_fn=eigval_fn, search_negative=search_negative, update_freq=update_freq)
-        super().__init__(defaults)
-
-        if inner is not None:
-            self.set_child('inner', inner)
-
-    @torch.no_grad
-    def update(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
-
-        settings = self.settings[params[0]]
-        damping = settings['damping']
-        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
-        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 = flatten_jacobian(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)
-
-            if damping != 0: H.add_(torch.eye(H.size(-1), dtype=H.dtype, device=H.device).mul_(damping))
-            self.global_state['H'] = H
+        defaults = locals().copy()
+        del defaults['self'], defaults['update_freq'], defaults["inner"]
+        super().__init__(defaults, update_freq=update_freq, inner=inner)
 
     @torch.no_grad
-    def apply(self, var):
-        H = self.global_state["H"]
-
-        params = var.params
-        settings = self.settings[params[0]]
-        search_negative = settings['search_negative']
-        H_tfm = settings['H_tfm']
-        eigval_fn = settings['eigval_fn']
-        use_lstsq = settings['use_lstsq']
-
-        # -------------------------------- 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])
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
 
-        # ----------------------------------- solve ---------------------------------- #
-        update = None
-        if H_tfm is not None:
-            ret = H_tfm(H, g)
+        _, _, self.global_state['H'] = objective.hessian(
+            hessian_method=fs['hessian_method'],
+            h=fs['h'],
+            at_x0=True
+        )
 
-            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_fn is not None):
-            update = _eigh_solve(H, g, eigval_fn, search_negative=search_negative)
-
-        if update is None and use_lstsq: update = _least_squares_solve(H, g)
-        if update is None: update = _cholesky_solve(H, g)
-        if update is None: update = _lu_solve(H, g)
-        if update is None: update = _least_squares_solve(H, g)
-
-        var.update = vec_to_tensors(update, params)
-
-        return var
-
-    def get_H(self,var):
-        H = self.global_state["H"]
-        settings = self.defaults
-        if settings['eigval_fn'] is not None:
-            try:
-                L, Q = torch.linalg.eigh(H) # pylint:disable=not-callable
-                L = settings['eigval_fn'](L)
-                H = Q @ L.diag_embed() @ Q.mH
-                H_inv = Q @ L.reciprocal().diag_embed() @ Q.mH
-                return DenseWithInverse(H, H_inv)
-
-            except torch.linalg.LinAlgError:
-                pass
-
-        return Dense(H)
-
-
-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).
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        params = objective.params
+        fs = settings[0]
 
-    .. warning::
-        this uses roughly O(N^2) memory.
+        update = _newton_step(
+            objective=objective,
+            H = self.global_state["H"],
+            damping = fs["damping"],
+            H_tfm = fs["H_tfm"],
+            eigval_fn = fs["eigval_fn"],
+            use_lstsq = fs["use_lstsq"],
+        )
 
-    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)
+        objective.updates = vec_to_tensors(update, params)
+        return objective
 
-        if inner is not None:
-            self.set_child('inner', inner)
+    def get_H(self,objective=...):
+        return _get_H(self.global_state["H"], self.defaults["eigval_fn"])
 
-    @torch.no_grad
-    def update(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 = flatten_jacobian(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)
-
-            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
-                finfo = torch.finfo(H.dtype)
-                Y = self.global_state['Y'] = num.div_(denom.clip(min=finfo.tiny * 2, max=finfo.max / 2))
-
-            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
-
-
-    def apply(self, var):
-        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
-
-    def get_H(self,var):
-        return DenseWithInverse(A = self.global_state["H"], A_inv=self.global_state["Y"])