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

torchzero/modules/experimental/higher_order_newton.py

@@ -1,21 +1,12 @@
-import itertools
 import math
-import warnings
-from collections.abc import Callable
-from contextlib import nullcontext
-from functools import partial
 from typing import Any, Literal
 
 import numpy as np
 import scipy.optimize
 import torch
 
-from ...core import Chainable, Module, apply_transform
+from ...core import DerivativesMethod, Module
 from ...utils import TensorList, vec_to_tensors, vec_to_tensors_
-from ...utils.derivatives import (
-    flatten_jacobian,
-    jacobian_wrt,
-)
 
 _LETTERS = 'abcdefghijklmnopqrstuvwxyz'
 def _poly_eval(s: np.ndarray, c, derivatives):
@@ -195,22 +186,22 @@ class HigherOrderNewton(Module):
         max_attempts = 10,
         boundary_tol: float = 1e-2,
         de_iters: int | None = None,
-        vectorize: bool = True,
+        derivatives_method: DerivativesMethod = "batched_autograd",
     ):
         if init is None:
             if trust_method == 'bounds': init = 1
             else: init = 0.1
 
-        defaults = dict(order=order, trust_method=trust_method, nplus=nplus, nminus=nminus, eta=eta, init=init, vectorize=vectorize, de_iters=de_iters, max_attempts=max_attempts, boundary_tol=boundary_tol, rho_good=rho_good, rho_bad=rho_bad)
+        defaults = dict(order=order, trust_method=trust_method, nplus=nplus, nminus=nminus, eta=eta, init=init, de_iters=de_iters, max_attempts=max_attempts, boundary_tol=boundary_tol, rho_good=rho_good, rho_bad=rho_bad, derivatives_method=derivatives_method)
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
+    def apply(self, objective):
+        params = TensorList(objective.params)
+        closure = objective.closure
         if closure is None: raise RuntimeError('HigherOrderNewton requires closure')
 
-        settings = self.settings[params[0]]
+        settings = self.defaults
         order = settings['order']
         nplus = settings['nplus']
         nminus = settings['nminus']
@@ -219,31 +210,12 @@ class HigherOrderNewton(Module):
         trust_method = settings['trust_method']
         de_iters = settings['de_iters']
         max_attempts = settings['max_attempts']
-        vectorize = settings['vectorize']
         boundary_tol = settings['boundary_tol']
         rho_good = settings['rho_good']
         rho_bad = settings['rho_bad']
 
         # ------------------------ calculate grad and hessian ------------------------ #
-        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()
-            derivatives = [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:])
-                    derivatives.append(T)
+        loss, *derivatives = objective.derivatives(order=order, at_x0=True, method=self.defaults["derivatives_method"])
 
         x0 = torch.cat([p.ravel() for p in params])
 
@@ -301,7 +273,8 @@ class HigherOrderNewton(Module):
                 vec_to_tensors_(x0, params)
                 reduction = loss - loss_star
 
-                rho = reduction / (max(pred_reduction, 1e-8))
+                rho = reduction / (max(pred_reduction, finfo.tiny * 2)) # pyright:ignore[reportArgumentType]
+
                 # failed step
                 if rho < rho_bad:
                     self.global_state['trust_region'] = trust_value * nminus
@@ -320,8 +293,9 @@ class HigherOrderNewton(Module):
         assert x_star is not None
         if success:
             difference = vec_to_tensors(x0 - x_star, params)
-            var.update = list(difference)
+            objective.updates = list(difference)
         else:
-            var.update = params.zeros_like()
-        return var
+            objective.updates = params.zeros_like()
+
+        return objective
 

torchzero/modules/experimental/l_infinity.py

@@ -43,7 +43,7 @@ class InfinityNormTrustRegion(TrustRegionBase):
 
         .. code-block:: python
 
-            opt = tz.Modular(
+            opt = tz.Optimizer(
                 model.parameters(),
                 tz.m.InfinityNormTrustRegion(hess_module=tz.m.BFGS(inverse=False)),
             )

torchzero/modules/experimental/matrix_nag.py

@@ -0,0 +1,122 @@
+from collections.abc import Callable
+from typing import Literal
+
+import torch
+from torchzero.core import Chainable, Transform, HVPMethod
+from torchzero.utils import NumberList, TensorList
+
+
+def matrix_nag_(
+    tensors_: TensorList,
+    s: TensorList,
+    Hvp_fn: Callable,
+    mu: float | NumberList,
+):
+    s += tensors_
+    Hv = TensorList(Hvp_fn(s))
+    s -= Hv.mul_(mu)
+    return tensors_.add_(s)
+
+
+class MatrixNAG(Transform):
+    """nesterov momentum version of matrix momentum. It seemed to work really well but adapting doesn't work,
+    I need to test more"""
+    def __init__(
+        self,
+        mu=0.1,
+        hvp_method: HVPMethod = "autograd",
+        h: float = 1e-3,
+        adaptive:bool = False,
+        adapt_freq: int | None = None,
+        hvp_tfm: Chainable | None = None,
+    ):
+        defaults = dict(mu=mu, hvp_method=hvp_method, h=h, adaptive=adaptive, adapt_freq=adapt_freq)
+        super().__init__(defaults)
+
+        if hvp_tfm is not None:
+            self.set_child('hvp_tfm', hvp_tfm)
+
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('p_prev')
+
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        assert objective.closure is not None
+        step = self.global_state.get("step", 0)
+        self.global_state["step"] = step + 1
+
+        p = TensorList(objective.params)
+        g = TensorList(objective.get_grads(create_graph=self.defaults["hvp_method"] == "autograd"))
+        p_prev = self.get_state(p, "p_prev", init=p, cls=TensorList)
+        s = p - p_prev
+        p_prev.copy_(p)
+
+        # -------------------------------- adaptive mu ------------------------------- #
+        if self.defaults["adaptive"]:
+
+            if step == 1:
+                self.global_state["mu_mul"] = 0
+
+            else:
+                # ---------------------------- deterministic case ---------------------------- #
+                if self.defaults["adapt_freq"] is None:
+                    g_prev = self.get_state(objective.params, "g_prev", cls=TensorList)
+                    y = g - g_prev
+                    g_prev.copy_(g)
+
+                    denom =  y.global_vector_norm()
+                    denom = denom.clip(min = torch.finfo(denom.dtype).tiny * 2)
+                    self.global_state["mu_mul"] = s.global_vector_norm() / denom
+
+                # -------------------------------- stochastic -------------------------------- #
+                else:
+                    adapt_freq = self.defaults["adapt_freq"]
+
+                    # we start on 1nd step, and want to adapt when we start, so use (step - 1)
+                    if (step - 1) % adapt_freq == 0:
+                        assert objective.closure is not None
+                        p_cur = p.clone()
+
+                        # move to previous params and evaluate p_prev with current mini-batch
+                        p.copy_(self.get_state(objective.params, 'p_prev'))
+                        with torch.enable_grad():
+                            objective.closure()
+                        g_prev = [t.grad if t.grad is not None else torch.zeros_like(t) for t in p]
+                        y = g - g_prev
+
+                        # move back to current params
+                        p.copy_(p_cur)
+
+                        denom =  y.global_vector_norm()
+                        denom = denom.clip(min = torch.finfo(denom.dtype).tiny * 2)
+                        self.global_state["mu_mul"] = s.global_vector_norm() / denom
+
+        # -------------------------- matrix momentum update -------------------------- #
+        mu = self.get_settings(p, "mu", cls=NumberList)
+        if "mu_mul" in self.global_state:
+            mu = mu * self.global_state["mu_mul"]
+
+        # def Hvp_fn(v):
+        #     Hv, _ = self.Hvp(
+        #         v=v,
+        #         at_x0=True,
+        #         var=objective,
+        #         rgrad=g,
+        #         hvp_method=self.defaults["hvp_method"],
+        #         h=self.defaults["h"],
+        #         normalize=True,
+        #         retain_grad=False,
+        #     )
+        #     return Hv
+
+        _, Hvp_fn = objective.list_Hvp_function(hvp_method=self.defaults["hvp_method"], h=self.defaults["h"], at_x0=True)
+
+        objective.updates = matrix_nag_(
+            tensors_=TensorList(objective.get_updates()),
+            s=s,
+            Hvp_fn=Hvp_fn,
+            mu=mu,
+        )
+
+        return objective

torchzero/modules/experimental/newton_solver.py

@@ -1,11 +1,10 @@
-from collections.abc import Callable, Iterable
-from typing import Any, Literal, overload
+from collections.abc import Callable
+from typing import Any
 
 import torch
 
-from ...core import Chainable, Modular, Module, apply_transform
-from ...utils import TensorList, as_tensorlist
-from ...utils.derivatives import hvp, hvp_fd_forward, hvp_fd_central
+from ...core import Chainable, Optimizer, Module, step, HVPMethod
+from ...utils import TensorList
 from ..quasi_newton import LBFGS
 
 
@@ -13,30 +12,32 @@ class NewtonSolver(Module):
     """Matrix free newton via with any custom solver (this is for testing, use NewtonCG or NystromPCG)."""
     def __init__(
         self,
-        solver: Callable[[list[torch.Tensor]], Any] = lambda p: Modular(p, LBFGS()),
+        solver: Callable[[list[torch.Tensor]], Any] = lambda p: Optimizer(p, LBFGS()),
         maxiter=None,
         maxiter1=None,
         tol:float | None=1e-3,
         reg: float = 0,
         warm_start=True,
-        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        hvp_method: HVPMethod = "autograd",
         reset_solver: bool = False,
         h: float= 1e-3,
+
         inner: Chainable | None = None,
     ):
-        defaults = dict(tol=tol, h=h,reset_solver=reset_solver, maxiter=maxiter, maxiter1=maxiter1, reg=reg, warm_start=warm_start, solver=solver, hvp_method=hvp_method)
-        super().__init__(defaults,)
+        defaults = locals().copy()
+        del defaults['self'], defaults['inner']
+        super().__init__(defaults)
 
-        if inner is not None:
-            self.set_child('inner', inner)
+        self.set_child("inner", inner)
 
         self._num_hvps = 0
         self._num_hvps_last_step = 0
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
+    def apply(self, objective):
+
+        params = TensorList(objective.params)
+        closure = objective.closure
         if closure is None: raise RuntimeError('NewtonCG requires closure')
 
         settings = self.settings[params[0]]
@@ -44,51 +45,19 @@ class NewtonSolver(Module):
         maxiter = settings['maxiter']
         maxiter1 = settings['maxiter1']
         tol = settings['tol']
-        reg = settings['reg']
         hvp_method = settings['hvp_method']
         warm_start = settings['warm_start']
         h = settings['h']
         reset_solver = settings['reset_solver']
 
         self._num_hvps_last_step = 0
-        # ---------------------- Hessian vector product function --------------------- #
-        if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
-
-            def H_mm(x):
-                self._num_hvps_last_step += 1
-                with torch.enable_grad():
-                    Hvp = TensorList(hvp(params, grad, x, retain_graph=True))
-                if reg != 0: Hvp = Hvp + (x*reg)
-                return Hvp
-
-        else:
-
-            with torch.enable_grad():
-                grad = var.get_grad()
-
-            if hvp_method == 'forward':
-                def H_mm(x):
-                    self._num_hvps_last_step += 1
-                    Hvp = TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1])
-                    if reg != 0: Hvp = Hvp + (x*reg)
-                    return Hvp
-
-            elif hvp_method == 'central':
-                def H_mm(x):
-                    self._num_hvps_last_step += 1
-                    Hvp =  TensorList(hvp_fd_central(closure, params, x, h=h, normalize=True)[1])
-                    if reg != 0: Hvp = Hvp + (x*reg)
-                    return Hvp
-
-            else:
-                raise ValueError(hvp_method)
 
+        # ---------------------- Hessian vector product function --------------------- #
+        _, H_mv = objective.list_Hvp_function(hvp_method=hvp_method, h=h, at_x0=True)
 
         # -------------------------------- inner step -------------------------------- #
-        b = as_tensorlist(grad)
-        if 'inner' in self.children:
-            b = as_tensorlist(apply_transform(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, var=var))
+        objective = self.inner_step("inner", objective, must_exist=False)
+        b = TensorList(objective.get_updates())
 
         # ---------------------------------- run cg ---------------------------------- #
         x0 = None
@@ -112,7 +81,7 @@ class NewtonSolver(Module):
                 solver = self.global_state['solver']
 
         def lstsq_closure(backward=True):
-            Hx = H_mm(x).detach()
+            Hx = H_mv(x).detach()
             # loss = (Hx-b).pow(2).global_mean()
             # if backward:
             #     solver.zero_grad()
@@ -122,7 +91,7 @@ class NewtonSolver(Module):
             loss = residual.pow(2).global_mean()
             if backward:
                 with torch.no_grad():
-                    H_residual = H_mm(residual)
+                    H_residual = H_mv(residual)
                     n = residual.global_numel()
                     x.set_grad_((2.0 / n) * H_residual)
 
@@ -143,8 +112,8 @@ class NewtonSolver(Module):
             assert x0 is not None
             x0.copy_(x)
 
-        var.update = x.detach()
+        objective.updates = x.detach()
         self._num_hvps += self._num_hvps_last_step
-        return var
+        return objective
 
 

torchzero/modules/experimental/newtonnewton.py

@@ -7,21 +7,21 @@ from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, apply_transform
-from ...utils import TensorList, vec_to_tensors
+from ...core import Chainable, Transform, step
+from ...linalg.linear_operator import Dense
+from ...utils import TensorList, vec_to_tensors_
 from ...utils.derivatives import (
     flatten_jacobian,
     jacobian_wrt,
 )
 from ..second_order.newton import (
-    _cholesky_solve,
-    _eigh_solve,
+    _try_cholesky_solve,
     _least_squares_solve,
-    _lu_solve,
+    _try_lu_solve,
 )
-from ...utils.linalg.linear_operator import Dense
 
-class NewtonNewton(Module):
+
+class NewtonNewton(Transform):
     """Applies Newton-like preconditioning to Newton step.
 
     This is a method that I thought of and then it worked. Here is how it works:
@@ -33,42 +33,36 @@ class NewtonNewton(Module):
     3. Solve H2 x2 = x for x2.
 
     4. Optionally, repeat (if order is higher than 3.)
-
-    Memory is n^order. It tends to converge faster on convex functions, but can be unstable on non-convex. Orders higher than 3 are usually too unsable and have little benefit.
-
-    3rd order variant can minimize some convex functions with up to 100 variables in less time than Newton's method,
-    this is if pytorch can vectorize hessian computation efficiently.
     """
     def __init__(
         self,
         reg: float = 1e-6,
         order: int = 3,
-        search_negative: bool = False,
         vectorize: bool = True,
-        eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        update_freq: int = 1,
+        inner: Chainable | None = None,
     ):
-        defaults = dict(order=order, reg=reg, vectorize=vectorize, eigval_fn=eigval_fn, search_negative=search_negative)
-        super().__init__(defaults)
+        defaults = dict(order=order, reg=reg, vectorize=vectorize)
+        super().__init__(defaults, update_freq=update_freq, inner=inner)
 
     @torch.no_grad
-    def update(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
+
+        params = TensorList(objective.params)
+        closure = objective.closure
         if closure is None: raise RuntimeError('NewtonNewton requires closure')
 
-        settings = self.settings[params[0]]
-        reg = settings['reg']
-        vectorize = settings['vectorize']
-        order = settings['order']
-        search_negative = settings['search_negative']
-        eigval_fn = settings['eigval_fn']
+        reg = fs['reg']
+        vectorize = fs['vectorize']
+        order = fs['order']
 
         # ------------------------ calculate grad and hessian ------------------------ #
-        Hs = []
+        P = None
         with torch.enable_grad():
-            loss = var.loss = var.loss_approx = closure(False)
+            loss = objective.loss = objective.loss_approx = closure(False)
             g_list = torch.autograd.grad(loss, params, create_graph=True)
-            var.grad = list(g_list)
+            objective.grads = list(g_list)
 
             xp = torch.cat([t.ravel() for t in g_list])
             I = torch.eye(xp.numel(), dtype=xp.dtype, device=xp.device)
@@ -79,27 +73,30 @@ class NewtonNewton(Module):
                 with torch.no_grad() if is_last else nullcontext():
                     H = flatten_jacobian(H_list)
                     if reg != 0: H = H + I * reg
-                    Hs.append(H)
+                    if P is None: P = H
+                    else: P = P @ H
+
+                    if not is_last:
+                        x = _try_cholesky_solve(H, xp)
+                        if x is None: x = _try_lu_solve(H, xp)
+                        if x is None: x = _least_squares_solve(H, xp)
+                        xp = x.squeeze()
+
+        self.global_state["P"] = P
+
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        updates = objective.get_updates()
+        P = self.global_state['P']
+        b = torch.cat([t.ravel() for t in updates])
 
-                    x = None
-                    if search_negative or (is_last and eigval_fn is not None):
-                        x = _eigh_solve(H, xp, eigval_fn, search_negative=search_negative)
-                    if x is None: x = _cholesky_solve(H, xp)
-                    if x is None: x = _lu_solve(H, xp)
-                    if x is None: x = _least_squares_solve(H, xp)
-                    xp = x.squeeze()
+        sol = _try_cholesky_solve(P, b)
+        if sol is None: sol = _try_lu_solve(P, b)
+        if sol is None: sol = _least_squares_solve(P, b)
 
-        self.global_state["Hs"] = Hs
-        self.global_state['xp'] = xp.nan_to_num_(0,0,0)
+        vec_to_tensors_(sol, updates)
+        return objective
 
     @torch.no_grad
-    def apply(self, var):
-        params = var.params
-        xp = self.global_state['xp']
-        var.update = vec_to_tensors(xp, params)
-        return var
-
-    def get_H(self, var):
-        Hs = self.global_state["Hs"]
-        if len(Hs) == 1: return Dense(Hs[0])
-        return Dense(torch.linalg.multi_dot(self.global_state["Hs"])) # pylint:disable=not-callable
+    def get_H(self, objective=...):
+        return Dense(self.global_state["P"])

torchzero/modules/experimental/reduce_outward_lr.py

@@ -1,28 +1,28 @@
 import torch
 
-from ...core import Target, Transform
+from ...core import  TensorTransform
 from ...utils import TensorList, unpack_states, unpack_dicts
 
-class ReduceOutwardLR(Transform):
+class ReduceOutwardLR(TensorTransform):
     """When update sign matches weight sign, the learning rate for that weight is multiplied by `mul`.
 
     This means updates that move weights towards zero have higher learning rates.
 
-    .. warning::
+    Warning:
         This sounded good but after testing turns out it sucks.
     """
-    def __init__(self, mul = 0.5, use_grad=False, invert=False, target: Target = 'update'):
+    def __init__(self, mul = 0.5, use_grad=False, invert=False):
         defaults = dict(mul=mul, use_grad=use_grad, invert=invert)
-        super().__init__(defaults, uses_grad=use_grad, target=target)
+        super().__init__(defaults, uses_grad=use_grad)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         params = TensorList(params)
         tensors = TensorList(tensors)
 
         mul = [s['mul'] for s in settings]
         s = settings[0]
-        use_grad = s['use_grad']
+        use_grad = self._uses_grad
         invert = s['invert']
 
         if use_grad: cur = grads

torchzero/modules/experimental/scipy_newton_cg.py

@@ -3,10 +3,9 @@ from typing import Literal, overload
 import torch
 from scipy.sparse.linalg import LinearOperator, gcrotmk
 
-from ...core import Chainable, Module, apply_transform
-from ...utils import NumberList, TensorList, as_tensorlist, generic_vector_norm, vec_to_tensors
-from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
-from ...utils.linalg.solve import cg, minres
+from ...core import Chainable, Module, step
+from ...utils import TensorList, vec_to_tensors
+from ...utils.derivatives import hvp_fd_central, hvp_fd_forward
 
 
 class ScipyNewtonCG(Module):
@@ -14,7 +13,7 @@ class ScipyNewtonCG(Module):
     def __init__(
         self,
         solver = gcrotmk,
-        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        hvp_method: Literal["fd_forward", "fd_central", "autograd"] = "autograd",
         h: float = 1e-3,
         warm_start=False,
         inner: Chainable | None = None,
@@ -33,47 +32,47 @@ class ScipyNewtonCG(Module):
         self._kwargs = kwargs
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
+    def apply(self, objective):
+        params = TensorList(objective.params)
+        closure = objective.closure
         if closure is None: raise RuntimeError('NewtonCG requires closure')
 
-        settings = self.settings[params[0]]
-        hvp_method = settings['hvp_method']
-        solver = settings['solver']
-        h = settings['h']
-        warm_start = settings['warm_start']
+        fs = self.settings[params[0]]
+        hvp_method = fs['hvp_method']
+        solver = fs['solver']
+        h = fs['h']
+        warm_start = fs['warm_start']
 
         self._num_hvps_last_step = 0
         # ---------------------- Hessian vector product function --------------------- #
         device = params[0].device; dtype=params[0].dtype
         if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
+            grad = objective.get_grads(create_graph=True)
 
             def H_mm(x_np):
                 self._num_hvps_last_step += 1
                 x = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), grad)
                 with torch.enable_grad():
-                    Hvp = TensorList(hvp(params, grad, x, retain_graph=True))
+                    Hvp = TensorList(torch.autograd.grad(grad, params, x, retain_graph=True))
                 return torch.cat([t.ravel() for t in Hvp]).numpy(force=True)
 
         else:
 
             with torch.enable_grad():
-                grad = var.get_grad()
+                grad = objective.get_grads()
 
             if hvp_method == 'forward':
                 def H_mm(x_np):
                     self._num_hvps_last_step += 1
                     x = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), grad)
-                    Hvp = TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1])
+                    Hvp = TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad)[1])
                     return torch.cat([t.ravel() for t in Hvp]).numpy(force=True)
 
             elif hvp_method == 'central':
                 def H_mm(x_np):
                     self._num_hvps_last_step += 1
                     x = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), grad)
-                    Hvp = TensorList(hvp_fd_central(closure, params, x, h=h, normalize=True)[1])
+                    Hvp = TensorList(hvp_fd_central(closure, params, x, h=h)[1])
                     return torch.cat([t.ravel() for t in Hvp]).numpy(force=True)
 
             else:
@@ -83,10 +82,8 @@ class ScipyNewtonCG(Module):
         H = LinearOperator(shape=(ndim,ndim), matvec=H_mm, rmatvec=H_mm) # type:ignore
 
         # -------------------------------- 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)
+        objective = self.inner_step("inner", objective, must_exist=False)
+        b = TensorList(objective.get_updates())
 
         # ---------------------------------- run cg ---------------------------------- #
         x0 = None
@@ -98,8 +95,8 @@ class ScipyNewtonCG(Module):
         if warm_start:
             self.global_state['x_prev'] = x_np
 
-        var.update = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), params)
+        objective.updates = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), params)
 
         self._num_hvps += self._num_hvps_last_step
-        return var
+        return objective