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

torchzero/modules/quasi_newton/quasi_newton.py

@@ -1,11 +1,13 @@
 """Use BFGS or maybe SR1."""
-from typing import Any, Literal
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
+from typing import Any, Literal
+
 import torch
 
-from ...core import Chainable, Module, Preconditioner, TensorwisePreconditioner
-from ...utils import TensorList, set_storage_
+from ...core import Chainable, Module, TensorwiseTransform, Transform
+from ...utils import TensorList, set_storage_, unpack_states
+
 
 def _safe_dict_update_(d1_:dict, d2:dict):
     inter = set(d1_.keys()).intersection(d2.keys())
@@ -17,14 +19,14 @@ def _maybe_lerp_(state, key, value: torch.Tensor, beta: float | None):
     elif state[key].shape != value.shape: state[key] = value
     else: state[key].lerp_(value, 1-beta)
 
-class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
+class HessianUpdateStrategy(TensorwiseTransform, ABC):
     def __init__(
         self,
         defaults: dict | None = None,
         init_scale: float | Literal["auto"] = "auto",
         tol: float = 1e-10,
         tol_reset: bool = True,
-        reset_interval: int | None = None,
+        reset_interval: int | None | Literal['auto'] = None,
         beta: float | None = None,
         update_freq: int = 1,
         scale_first: bool = True,
@@ -44,7 +46,7 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
         if ys != 0 and yy != 0: return yy/ys
         return 1
 
-    def _reset_M_(self, M: torch.Tensor, s:torch.Tensor,y:torch.Tensor,inverse:bool, init_scale: Any):
+    def _reset_M_(self, M: torch.Tensor, s:torch.Tensor,y:torch.Tensor, inverse:bool, init_scale: Any, state:dict[str,Any]):
         set_storage_(M, torch.eye(M.size(-1), device=M.device, dtype=M.dtype))
         if init_scale == 'auto': init_scale = self._get_init_scale(s,y)
         if init_scale >= 1:
@@ -62,7 +64,7 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
         raise NotImplementedError
 
     @torch.no_grad
-    def update_tensor(self, tensor, param, grad, state, settings):
+    def update_tensor(self, tensor, param, grad, loss, state, settings):
         p = param.view(-1); g = tensor.view(-1)
         inverse = settings['inverse']
         M_key = 'H' if inverse else 'B'
@@ -73,6 +75,7 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
         tol = settings['tol']
         tol_reset = settings['tol_reset']
         reset_interval = settings['reset_interval']
+        if reset_interval == 'auto': reset_interval = tensor.numel() + 1
 
         if M is None:
             M = torch.eye(p.size(0), device=p.device, dtype=p.dtype)
@@ -81,10 +84,12 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
                 else: M *= init_scale
 
             state[M_key] = M
+            state['f_prev'] = loss
             state['p_prev'] = p.clone()
             state['g_prev'] = g.clone()
             return
 
+        state['f'] = loss
         p_prev = state['p_prev']
         g_prev = state['g_prev']
         s: torch.Tensor = p - p_prev
@@ -93,13 +98,13 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
         state['g_prev'].copy_(g)
 
         if reset_interval is not None and step != 0 and step % reset_interval == 0:
-            self._reset_M_(M, s, y, inverse, init_scale)
+            self._reset_M_(M, s, y, inverse, init_scale, state)
             return
 
         # tolerance on gradient difference to avoid exploding after converging
-        elif y.abs().max() <= tol:
+        if y.abs().max() <= tol:
             # reset history
-            if tol_reset: self._reset_M_(M, s, y, inverse, init_scale)
+            if tol_reset: self._reset_M_(M, s, y, inverse, init_scale, state)
             return
 
         if step == 1 and init_scale == 'auto':
@@ -117,8 +122,10 @@ class HessianUpdateStrategy(TensorwisePreconditioner, ABC):
             B_new = self.update_B(B=M, s=s, y=y, p=p, g=g, p_prev=p_prev, g_prev=g_prev, state=state, settings=settings)
             _maybe_lerp_(state, 'B', B_new, beta)
 
+        state['f_prev'] = loss
+
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, state, settings):
+    def apply_tensor(self, tensor, param, grad, loss, state, settings):
         step = state.get('step', 0)
 
         if settings['scale_second'] and step == 2:
@@ -198,19 +205,15 @@ class SR1(HUpdateStrategy):
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
         return sr1_H_(H=H, s=s, y=y, tol=settings['tol'])
 
-# BFGS has defaults - init_scale = "auto" and scale_second = False
-# SR1 has defaults -  init_scale = 1 and scale_second = True
-# basically some methods work better with first and some with second.
-# I inherit from BFGS or SR1 to avoid writing all those arguments again
 # ------------------------------------ DFP ----------------------------------- #
 def dfp_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
     sy = torch.dot(s, y)
     if sy.abs() <= tol: return H
     term1 = torch.outer(s, s).div_(sy)
-    denom = torch.dot(y, H @ y) #
-    if denom.abs() <= tol: return H
+    yHy = torch.dot(y, H @ y) #
+    if yHy.abs() <= tol: return H
     num = H @ torch.outer(y, y) @ H
-    term2 = num.div_(denom)
+    term2 = num.div_(yHy)
     H += term1.sub_(term2)
     return H
 
@@ -225,34 +228,35 @@ class DFP(HUpdateStrategy):
 
 def broyden_good_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
     c = H.T @ s
-    denom = c.dot(y)
-    if denom.abs() <= tol: return H
+    cy = c.dot(y)
+    if cy.abs() <= tol: return H
     num = (H@y).sub_(s).outer(c)
-    H -= num/denom
+    H -= num/cy
     return H
 
 def broyden_bad_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
     c = y
-    denom = c.dot(y)
-    if denom.abs() <= tol: return H
+    cy = c.dot(y)
+    if cy.abs() <= tol: return H
     num = (H@y).sub_(s).outer(c)
-    H -= num/denom
+    H -= num/cy
     return H
 
 def greenstadt1_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, g_prev: torch.Tensor, tol: float):
     c = g_prev
-    denom = c.dot(y)
-    if denom.abs() <= tol: return H
+    cy = c.dot(y)
+    if cy.abs() <= tol: return H
     num = (H@y).sub_(s).outer(c)
-    H -= num/denom
+    H -= num/cy
     return H
 
 def greenstadt2_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
-    c = torch.linalg.multi_dot([H,H,y]) # pylint:disable=not-callable
-    denom = c.dot(y)
-    if denom.abs() <= tol: return H
-    num = (H@y).sub_(s).outer(c)
-    H -= num/denom
+    Hy = H @ y
+    c = H @ Hy # pylint:disable=not-callable
+    cy = c.dot(y)
+    if cy.abs() <= tol: return H
+    num = Hy.sub_(s).outer(c)
+    H -= num/cy
     return H
 
 class BroydenGood(HUpdateStrategy):
@@ -273,11 +277,7 @@ class Greenstadt2(HUpdateStrategy):
 
 
 def column_updating_H_(H:torch.Tensor, s:torch.Tensor, y:torch.Tensor, tol:float):
-    n = H.shape[0]
-
     j = y.abs().argmax()
-    u = torch.zeros(n, device=H.device, dtype=H.dtype)
-    u[j] = 1.0
 
     denom = y[j]
     if denom.abs() < tol: return H
@@ -297,15 +297,15 @@ def thomas_H_(H: torch.Tensor, R:torch.Tensor, s: torch.Tensor, y: torch.Tensor,
     s_norm = torch.linalg.vector_norm(s) # pylint:disable=not-callable
     I = torch.eye(H.size(-1), device=H.device, dtype=H.dtype)
     d = (R + I * (s_norm/2)) @ s
-    denom = d.dot(s)
-    if denom.abs() <= tol: return H, R
-    R = (1 + s_norm) * ((I*s_norm).add_(R).sub_(d.outer(d).div_(denom)))
+    ds = d.dot(s)
+    if ds.abs() <= tol: return H, R
+    R = (1 + s_norm) * ((I*s_norm).add_(R).sub_(d.outer(d).div_(ds)))
 
     c = H.T @ d
-    denom = c.dot(y)
-    if denom.abs() <= tol: return H, R
+    cy = c.dot(y)
+    if cy.abs() <= tol: return H, R
     num = (H@y).sub_(s).outer(c)
-    H -= num/denom
+    H -= num/cy
     return H, R
 
 class ThomasOptimalMethod(HUpdateStrategy):
@@ -315,6 +315,11 @@ class ThomasOptimalMethod(HUpdateStrategy):
         H, state['R'] = thomas_H_(H=H, R=state['R'], s=s, y=y, tol=settings['tol'])
         return H
 
+    def _reset_M_(self, M, s, y,inverse, init_scale, state):
+        super()._reset_M_(M, s, y, inverse, init_scale, state)
+        for st in self.state.values():
+            st.pop("R", None)
+
 # ------------------------ powell's symmetric broyden ------------------------ #
 def psb_B_(B: torch.Tensor, s: torch.Tensor, y: torch.Tensor, tol:float):
     y_Bs = y - B@s
@@ -326,6 +331,7 @@ def psb_B_(B: torch.Tensor, s: torch.Tensor, y: torch.Tensor, tol:float):
     B += term1.sub_(term2)
     return B
 
+# I couldn't find formula for H
 class PSB(HessianUpdateStrategy):
     def __init__(
         self,
@@ -358,17 +364,85 @@ class PSB(HessianUpdateStrategy):
     def update_B(self, B, s, y, p, g, p_prev, g_prev, state, settings):
         return psb_B_(B=B, s=s, y=y, tol=settings['tol'])
 
-def pearson2_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
+
+# Algorithms from Pearson, J. D. (1969). Variable metric methods of minimisation. The Computer Journal, 12(2), 171–178. doi:10.1093/comjnl/12.2.171
+def pearson_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
+    Hy = H@y
+    yHy = y.dot(Hy)
+    if yHy.abs() <= tol: return H
+    num = (s - Hy).outer(Hy)
+    H += num.div_(yHy)
+    return H
+
+class Pearson(HUpdateStrategy):
+    """Pearson, J. D. (1969). Variable metric methods of minimisation. The Computer Journal, 12(2), 171–178. doi:10.1093/comjnl/12.2.171.
+
+    This is "Algorithm 2", attributed to McCormick in this paper. However for some reason this method is also called Pearson's 2nd method."""
+    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
+        return pearson_H_(H=H, s=s, y=y, tol=settings['tol'])
+
+def mccormick_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
     sy = s.dot(y)
     if sy.abs() <= tol: return H
     num = (s - H@y).outer(s)
     H += num.div_(sy)
     return H
 
-class Pearson2(HUpdateStrategy):
-    """finally found a reference in https://www.recotechnologies.com/~beigi/ps/asme-jdsmc-93-2.pdf"""
+class McCormick(HUpdateStrategy):
+    """Pearson, J. D. (1969). Variable metric methods of minimisation. The Computer Journal, 12(2), 171–178. doi:10.1093/comjnl/12.2.171.
+
+    This is "Algorithm 2", attributed to McCormick in this paper. However for some reason this method is also called Pearson's 2nd method."""
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
-        return pearson2_H_(H=H, s=s, y=y, tol=settings['tol'])
+        return mccormick_H_(H=H, s=s, y=y, tol=settings['tol'])
+
+def projected_newton_raphson_H_(H: torch.Tensor, R:torch.Tensor, s: torch.Tensor, y: torch.Tensor, tol:float):
+    Hy = H @ y
+    yHy = y.dot(Hy)
+    if yHy.abs() < tol: return H, R
+    H -= Hy.outer(Hy) / yHy
+    R += (s - R@y).outer(Hy) / yHy
+    return H, R
+
+class ProjectedNewtonRaphson(HessianUpdateStrategy):
+    """Pearson, J. D. (1969). Variable metric methods of minimisation. The Computer Journal, 12(2), 171–178. doi:10.1093/comjnl/12.2.171.
+
+    Algorithm 7"""
+    def __init__(
+        self,
+        init_scale: float | Literal["auto"] = 'auto',
+        tol: float = 1e-10,
+        tol_reset: bool = True,
+        reset_interval: int | None | Literal['auto'] = 'auto',
+        beta: float | None = None,
+        update_freq: int = 1,
+        scale_first: bool = True,
+        scale_second: bool = False,
+        concat_params: bool = True,
+        inner: Chainable | None = None,
+    ):
+        super().__init__(
+            init_scale=init_scale,
+            tol=tol,
+            tol_reset=tol_reset,
+            reset_interval=reset_interval,
+            beta=beta,
+            update_freq=update_freq,
+            scale_first=scale_first,
+            scale_second=scale_second,
+            concat_params=concat_params,
+            inverse=True,
+            inner=inner,
+        )
+
+    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
+        if 'R' not in state: state['R'] = torch.eye(H.size(-1), device=H.device, dtype=H.dtype)
+        H, R = projected_newton_raphson_H_(H=H, R=state['R'], s=s, y=y, tol=settings['tol'])
+        state["R"] = R
+        return H
+
+    def _reset_M_(self, M, s, y, inverse, init_scale, state):
+        assert inverse
+        M.copy_(state["R"])
 
 # Oren, S. S., & Spedicato, E. (1976). Optimal conditioning of self-scaling variable metric algorithms. Mathematical programming, 10(1), 70-90.
 def ssvm_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, g:torch.Tensor, switch: tuple[float,float] | Literal[1,2,3,4], tol: float):
@@ -473,4 +547,137 @@ class SSVM(HessianUpdateStrategy):
         )
 
     def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
-        return ssvm_H_(H=H, s=s, y=y, g=g, switch=settings['switch'], tol=settings['tol'])
+        return ssvm_H_(H=H, s=s, y=y, g=g, switch=settings['switch'], tol=settings['tol'])
+
+# HOSHINO, S. (1972). A Formulation of Variable Metric Methods. IMA Journal of Applied Mathematics, 10(3), 394–403. doi:10.1093/imamat/10.3.394
+def hoshino_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
+    Hy = H@y
+    ys = y.dot(s)
+    if ys.abs() <= tol: return H
+    yHy = y.dot(Hy)
+    denom = ys + yHy
+    if denom.abs() <= tol: return H
+
+    term1 = 1/denom
+    term2 = s.outer(s).mul_(1 + ((2 * yHy) / ys))
+    term3 = s.outer(y) @ H
+    term4 = Hy.outer(s)
+    term5 = Hy.outer(y) @ H
+
+    inner_term = term2 - term3 - term4 - term5
+    H += inner_term.mul_(term1)
+    return H
+
+def gradient_correction(g: TensorList, s: TensorList, y: TensorList):
+    sy = s.dot(y)
+    if sy.abs() < torch.finfo(g[0].dtype).eps: return g
+    return g - (y * (s.dot(g) / sy))
+
+
+class GradientCorrection(Transform):
+    """estimates gradient at minima along search direction assuming function is quadratic as proposed in HOSHINO, S. (1972). A Formulation of Variable Metric Methods. IMA Journal of Applied Mathematics, 10(3), 394–403. doi:10.1093/imamat/10.3.394
+
+    This can useful as inner module for second order methods."""
+    def __init__(self):
+        super().__init__(None, uses_grad=False)
+
+    def apply(self, tensors, params, grads, loss, states, settings):
+        if 'p_prev' not in states[0]:
+            p_prev = unpack_states(states, tensors, 'p_prev', init=params)
+            g_prev = unpack_states(states, tensors, 'g_prev', init=tensors)
+            return tensors
+
+        p_prev, g_prev = unpack_states(states, tensors, 'p_prev', 'g_prev', cls=TensorList)
+        g_hat = gradient_correction(TensorList(tensors), params-p_prev, tensors-g_prev)
+
+        p_prev.copy_(params)
+        g_prev.copy_(tensors)
+        return g_hat
+
+class Horisho(HUpdateStrategy):
+    """HOSHINO, S. (1972). A Formulation of Variable Metric Methods. IMA Journal of Applied Mathematics, 10(3), 394–403. doi:10.1093/imamat/10.3.394"""
+    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
+        return hoshino_H_(H=H, s=s, y=y, tol=settings['tol'])
+
+# Fletcher, R. (1970). A new approach to variable metric algorithms. The Computer Journal, 13(3), 317–322. doi:10.1093/comjnl/13.3.317
+def fletcher_vmm_H_(H:torch.Tensor, s: torch.Tensor, y:torch.Tensor, tol: float):
+    sy = s.dot(y)
+    if sy.abs() < tol: return H
+    Hy = H @ y
+
+    term1 = (s.outer(y) @ H).div_(sy)
+    term2 = (Hy.outer(s)).div_(sy)
+    term3 = 1 + (y.dot(Hy) / sy)
+    term4 = s.outer(s).div_(sy)
+
+    H -= (term1 + term2 - term4.mul_(term3))
+    return H
+
+class FletcherVMM(HUpdateStrategy):
+    """Fletcher, R. (1970). A new approach to variable metric algorithms. The Computer Journal, 13(3), 317–322. doi:10.1093/comjnl/13.3.317"""
+    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
+        return fletcher_vmm_H_(H=H, s=s, y=y, tol=settings['tol'])
+
+
+# Moghrabi, I. A., Hassan, B. A., & Askar, A. (2022). New self-scaling quasi-newton methods for unconstrained optimization. Int. J. Math. Comput. Sci., 17, 1061U.
+def new_ssm1(H: torch.Tensor, s: torch.Tensor, y: torch.Tensor, f, f_prev, tol: float, type:int):
+    sy = s.dot(y)
+    if sy < tol: return H
+
+    term1 = (H @ y.outer(s) + s.outer(y) @ H) / sy
+
+    if type == 1:
+        pba = (2*sy + 2*(f-f_prev)) / sy
+
+    elif type == 2:
+        pba = (f_prev - f + 1/(2*sy)) / sy
+
+    else:
+        raise RuntimeError(type)
+
+    term3 = 1/pba + y.dot(H@y) / sy
+    term4 = s.outer(s) / sy
+
+    H.sub_(term1)
+    H.add_(term4.mul_(term3))
+    return H
+
+
+class NewSSM(HessianUpdateStrategy):
+    """Self-scaling method, requires a line search.
+
+    Moghrabi, I. A., Hassan, B. A., & Askar, A. (2022). New self-scaling quasi-newton methods for unconstrained optimization. Int. J. Math. Comput. Sci., 17, 1061U."""
+    def __init__(
+        self,
+        type: Literal[1, 2] = 1,
+        init_scale: float | Literal["auto"] = "auto",
+        tol: float = 1e-10,
+        tol_reset: bool = True,
+        reset_interval: int | None = None,
+        beta: float | None = None,
+        update_freq: int = 1,
+        scale_first: bool = True,
+        scale_second: bool = False,
+        concat_params: bool = True,
+        inner: Chainable | None = None,
+    ):
+        super().__init__(
+            defaults=dict(type=type),
+            init_scale=init_scale,
+            tol=tol,
+            tol_reset=tol_reset,
+            reset_interval=reset_interval,
+            beta=beta,
+            update_freq=update_freq,
+            scale_first=scale_first,
+            scale_second=scale_second,
+            concat_params=concat_params,
+            inverse=True,
+            inner=inner,
+        )
+    def update_H(self, H, s, y, p, g, p_prev, g_prev, state, settings):
+        f = state['f']
+        f_prev = state['f_prev']
+        return new_ssm1(H=H, s=s, y=y, f=f, f_prev=f_prev, type=settings['type'], tol=settings['tol'])
+
+

torchzero/modules/second_order/newton.py

@@ -5,7 +5,7 @@ from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, apply
+from ...core import Chainable, Module, apply_transform
 from ...utils import TensorList, vec_to_tensors
 from ...utils.derivatives import (
     hessian_list_to_mat,
@@ -18,9 +18,12 @@ from ...utils.derivatives import (
 
 
 def lu_solve(H: torch.Tensor, g: torch.Tensor):
-    x, info = torch.linalg.solve_ex(H, g) # pylint:disable=not-callable
-    if info == 0: return x
-    return None
+    try:
+        x, info = torch.linalg.solve_ex(H, g) # pylint:disable=not-callable
+        if info == 0: return x
+        return None
+    except RuntimeError:
+        return None
 
 def cholesky_solve(H: torch.Tensor, g: torch.Tensor):
     x, info = torch.linalg.cholesky_ex(H) # pylint:disable=not-callable
@@ -32,10 +35,15 @@ def cholesky_solve(H: torch.Tensor, g: torch.Tensor):
 def least_squares_solve(H: torch.Tensor, g: torch.Tensor):
     return torch.linalg.lstsq(H, g)[0] # pylint:disable=not-callable
 
-def eigh_solve(H: torch.Tensor, g: torch.Tensor, tfm: Callable | None):
+def eigh_solve(H: torch.Tensor, g: torch.Tensor, tfm: Callable | None, search_negative: bool):
     try:
         L, Q = torch.linalg.eigh(H) # pylint:disable=not-callable
         if tfm is not None: L = tfm(L)
+        if search_negative and L[0] < 0:
+            d = Q[0]
+             # 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
     except torch.linalg.LinAlgError:
@@ -56,6 +64,8 @@ class Newton(Module):
     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.
         hessian_method (str):
             how to calculate hessian. Defaults to "autograd".
         vectorize (bool, optional):
@@ -75,27 +85,29 @@ class Newton(Module):
         self,
         reg: float = 1e-6,
         eig_reg: bool = False,
+        search_negative: bool = False,
         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,
         eigval_tfm: Callable[[torch.Tensor], torch.Tensor] | None = None,
     ):
-        defaults = dict(reg=reg, eig_reg=eig_reg, abs=abs,hessian_method=hessian_method, vectorize=vectorize, H_tfm=H_tfm, eigval_tfm=eigval_tfm)
+        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)
         super().__init__(defaults)
 
         if inner is not None:
             self.set_child('inner', inner)
 
     @torch.no_grad
-    def step(self, vars):
-        params = TensorList(vars.params)
-        closure = vars.closure
+    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]]
         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']
@@ -104,16 +116,16 @@ class Newton(Module):
         # ------------------------ calculate grad and hessian ------------------------ #
         if hessian_method == 'autograd':
             with torch.enable_grad():
-                loss = vars.loss = vars.loss_approx = closure(False)
+                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
-                vars.grad = g_list
+                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 = vars.get_grad(retain_graph=True)
+                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]
 
@@ -121,10 +133,10 @@ class Newton(Module):
             raise ValueError(hessian_method)
 
         # -------------------------------- inner step -------------------------------- #
-        update = vars.get_update()
+        update = var.get_update()
         if 'inner' in self.children:
-            update = apply(self.children['inner'], update, params=params, grads=list(g_list), vars=vars)
-        g = torch.cat([t.view(-1) for t in update])
+            update = apply_transform(self.children['inner'], update, params=params, grads=list(g_list), var=var)
+        g = torch.cat([t.ravel() for t in update])
 
         # ------------------------------- regulazition ------------------------------- #
         if eig_reg: H = eig_tikhonov_(H, reg)
@@ -134,14 +146,14 @@ class Newton(Module):
         update = None
         if H_tfm is not None:
             H, is_inv = H_tfm(H, g)
-            if is_inv: update = H
+            if is_inv: update = H @ g
 
-        if eigval_tfm is not None:
-            update = eigh_solve(H, g, eigval_tfm)
+        if search_negative or (eigval_tfm is not None):
+            update = eigh_solve(H, g, eigval_tfm, search_negative=search_negative)
 
         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)
 
-        vars.update = vec_to_tensors(update, params)
-        return vars
+        var.update = vec_to_tensors(update, params)
+        return var

torchzero/modules/second_order/newton_cg.py

@@ -6,14 +6,14 @@ import torch
 from ...utils import TensorList, as_tensorlist, generic_zeros_like, generic_vector_norm, generic_numel
 from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
 
-from ...core import Chainable, apply, Module
+from ...core import Chainable, apply_transform, Module
 from ...utils.linalg.solve import cg
 
 class NewtonCG(Module):
     def __init__(
         self,
         maxiter=None,
-        tol=1e-3,
+        tol=1e-4,
         reg: float = 1e-8,
         hvp_method: Literal["forward", "central", "autograd"] = "forward",
         h=1e-3,
@@ -27,9 +27,9 @@ class NewtonCG(Module):
             self.set_child('inner', inner)
 
     @torch.no_grad
-    def step(self, vars):
-        params = TensorList(vars.params)
-        closure = vars.closure
+    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]]
@@ -42,7 +42,7 @@ class NewtonCG(Module):
 
         # ---------------------- Hessian vector product function --------------------- #
         if hvp_method == 'autograd':
-            grad = vars.get_grad(create_graph=True)
+            grad = var.get_grad(create_graph=True)
 
             def H_mm(x):
                 with torch.enable_grad():
@@ -51,7 +51,7 @@ class NewtonCG(Module):
         else:
 
             with torch.enable_grad():
-                grad = vars.get_grad()
+                grad = var.get_grad()
 
             if hvp_method == 'forward':
                 def H_mm(x):
@@ -66,19 +66,20 @@ class NewtonCG(Module):
 
 
         # -------------------------------- inner step -------------------------------- #
-        b = vars.get_update()
+        b = var.get_update()
         if 'inner' in self.children:
-            b = as_tensorlist(apply(self.children['inner'], b, params=params, grads=grad, vars=vars))
+            b = as_tensorlist(apply_transform(self.children['inner'], b, params=params, grads=grad, var=var))
 
         # ---------------------------------- run cg ---------------------------------- #
         x0 = None
-        if warm_start: x0 = self.get_state('prev_x', params=params, cls=TensorList) # initialized to 0 which is default anyway
+        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 warm_start:
             assert x0 is not None
             x0.copy_(x)
 
-        vars.update = x
-        return vars
+        var.update = x
+        return var