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

torchzero/modules/experimental/diagonal_higher_order_newton.py

@@ -1,225 +0,0 @@
-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 ...utils import TensorList, vec_to_tensors, vec_to_tensors_
-from ...utils.derivatives import (
-    hessian_list_to_mat,
-    jacobian_wrt,
-    hvp,
-)
-
-def _poly_eval_diag(s: np.ndarray, c, derivatives):
-    val = float(c) + (derivatives[0] * s).sum(-1)
-
-    if len(derivatives) > 1:
-        for i, d_diag in enumerate(derivatives[1:], 2):
-            val += (d_diag * (s**i)).sum(-1) / math.factorial(i)
-
-    return val
-
-def _proximal_poly_v_diag(x: np.ndarray, c, prox, x0: np.ndarray, derivatives):
-    """Computes the value of the proximal polynomial approximation."""
-    if x.ndim == 2: x = x.T
-    s = x - x0
-
-    val = _poly_eval_diag(s, c, derivatives)
-
-    penalty = 0
-    if prox != 0:
-        penalty = (prox / 2) * (s**2).sum(-1)
-
-    return val + penalty
-
-def _proximal_poly_g_diag(x: np.ndarray, c, prox, x0: np.ndarray, derivatives):
-    """Computes the gradient of the proximal polynomial approximation."""
-    s = x - x0
-
-    g = derivatives[0].copy()
-
-    if len(derivatives) > 1:
-        for i, d_diag in enumerate(derivatives[1:], 2):
-            g += d_diag * (s**(i - 1)) / math.factorial(i - 1)
-
-    if prox != 0:
-        g += prox * s
-
-    return g
-
-def _proximal_poly_H_diag(x: np.ndarray, c, prox, x0: np.ndarray, derivatives):
-    """Computes the Hessian of the proximal polynomial approximation."""
-    s = x - x0
-    n = x.shape[0]
-
-    if len(derivatives) < 2:
-        H_diag = np.zeros(n, dtype=s.dtype)
-    else:
-        H_diag = derivatives[1].copy()
-
-    if len(derivatives) > 2:
-        for i, d_diag in enumerate(derivatives[2:], 3):
-            H_diag += d_diag * (s**(i - 2)) / math.factorial(i - 2)
-
-    if prox != 0:
-        H_diag += prox
-
-    return np.diag(H_diag)
-
-def _poly_minimize(trust_region, prox, de_iters: Any, c, x: torch.Tensor, derivatives):
-    derivatives = [T.detach().cpu().numpy().astype(np.float64) for T in derivatives]
-    x0 = x.detach().cpu().numpy().astype(np.float64) # taylor series center
-    bounds = None
-    if trust_region is not None: bounds = list(zip(x0 - trust_region, x0 + trust_region))
-
-    # if len(derivatives) is 1, only gradient is available, I use that to test proximal penalty and bounds
-    if bounds is None:
-        if len(derivatives) == 1: method = 'bfgs'
-        else: method = 'trust-exact'
-    else:
-        if len(derivatives) == 1: method = 'l-bfgs-b'
-        else: method = 'trust-constr'
-
-    x_init = x0.copy()
-    v0 = _proximal_poly_v_diag(x0, c, prox, x0, derivatives)
-    if de_iters is not None and de_iters != 0:
-        if de_iters == -1: de_iters = None # let scipy decide
-        res = scipy.optimize.differential_evolution(
-            _proximal_poly_v_diag,
-            bounds if bounds is not None else list(zip(x0 - 10, x0 + 10)),
-            args=(c, prox, x0.copy(), derivatives),
-            maxiter=de_iters,
-            vectorized=True,
-        )
-        if res.fun < v0: x_init = res.x
-
-    res = scipy.optimize.minimize(
-        _proximal_poly_v_diag,
-        x_init,
-        method=method,
-        args=(c, prox, x0.copy(), derivatives),
-        jac=_proximal_poly_g_diag,
-        hess=_proximal_poly_H_diag,
-        bounds=bounds
-    )
-
-    return torch.from_numpy(res.x).to(x), res.fun
-
-
-
-class DiagonalHigherOrderNewton(Module):
-    """
-    Hvp with ones doesn't give you the diagonal unless derivatives are diagonal, but somehow it still works,
-    except it doesn't work in all cases except ones where it works.
-    """
-    def __init__(
-        self,
-        order: int = 4,
-        trust_method: Literal['bounds', 'proximal', 'none'] | None = 'bounds',
-        increase: float = 1.5,
-        decrease: float = 0.75,
-        trust_init: float | None = None,
-        trust_tol: float = 1,
-        de_iters: int | None = None,
-        vectorize: bool = True,
-    ):
-        if trust_init is None:
-            if trust_method == 'bounds': trust_init = 1
-            else: trust_init = 0.1
-
-        defaults = dict(order=order, trust_method=trust_method, increase=increase, decrease=decrease, trust_tol=trust_tol, trust_init=trust_init, vectorize=vectorize, de_iters=de_iters)
-        super().__init__(defaults)
-
-    @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]]
-        order = settings['order']
-        increase = settings['increase']
-        decrease = settings['decrease']
-        trust_tol = settings['trust_tol']
-        trust_init = settings['trust_init']
-        trust_method = settings['trust_method']
-        de_iters = settings['de_iters']
-
-        trust_value = self.global_state.get('trust_value', trust_init)
-
-
-        # ------------------------ calculate grad and hessian ------------------------ #
-        with torch.enable_grad():
-            loss = var.loss = var.loss_approx = closure(False)
-
-            g = torch.autograd.grad(loss, params, create_graph=True)
-            var.grad = list(g)
-
-            derivatives = [g]
-            T = g # current derivatives tensor diagonal
-            ones = [torch.ones_like(t) for t in g]
-
-            # get all derivatives up to order
-            for o in range(2, order + 1):
-                T = hvp(params, T, ones, create_graph=o != order)
-                derivatives.append(T)
-
-        x0 = torch.cat([p.ravel() for p in params])
-
-        if trust_method is None: trust_method = 'none'
-        else: trust_method = trust_method.lower()
-
-        if trust_method == 'none':
-            trust_region = None
-            prox = 0
-
-        elif trust_method == 'bounds':
-            trust_region = trust_value
-            prox = 0
-
-        elif trust_method == 'proximal':
-            trust_region = None
-            prox = 1 / trust_value
-
-        else:
-            raise ValueError(trust_method)
-
-        x_star, expected_loss = _poly_minimize(
-            trust_region=trust_region,
-            prox=prox,
-            de_iters=de_iters,
-            c=loss.item(),
-            x=x0,
-            derivatives=[torch.cat([t.ravel() for t in d]) for d in derivatives],
-        )
-
-        # trust region
-        if trust_method != 'none':
-            expected_reduction = loss - expected_loss
-
-            vec_to_tensors_(x_star, params)
-            loss_star = closure(False)
-            vec_to_tensors_(x0, params)
-            reduction = loss - loss_star
-
-            # failed step
-            if reduction <= 0:
-                x_star = x0
-                self.global_state['trust_value'] = trust_value * decrease
-
-            # very good step
-            elif expected_reduction / reduction <= trust_tol:
-                self.global_state['trust_value'] = trust_value * increase
-
-        difference = vec_to_tensors(x0 - x_star, params)
-        var.update = list(difference)
-        return var
-

torchzero/modules/experimental/soapy.py

@@ -1,163 +0,0 @@
-from operator import itemgetter
-
-import torch
-
-from ...core import Chainable, Transform
-from ..optimizers.shampoo import _merge_small_dims, _unmerge_small_dims
-from ..optimizers.soap import (
-    update_soap_covariances_,
-    get_orthogonal_matrix,
-    get_orthogonal_matrix_QR,
-    project,
-    project_back,
-)
-
-class SOAPY(Transform):
-    """Adam but uses scaled gradient differences for GGᵀ. Please note that this is experimental and isn't guaranteed to work.
-
-    New args:
-        scale_by_s - whether to scale gradient differences by parameter differences
-        y_to_ema2 - whether to use gradient differences for exponential moving average too
-    """
-    def __init__(
-        self,
-        beta1: float = 0.95,
-        beta2: float = 0.95,
-        shampoo_beta: float | None = 0.95,
-        precond_freq: int = 10,
-        merge_small: bool = True,
-        max_dim: int = 2_000,
-        precondition_1d: bool = True,
-        eps: float = 1e-8,
-        decay: float | None = None,
-        alpha: float = 1,
-        bias_correction: bool = True,
-        scale_by_s: bool = True,
-        y_to_ema2: bool = False,
-    ):
-        defaults = dict(
-            beta1=beta1,
-            beta2=beta2,
-            shampoo_beta=shampoo_beta,
-            precond_freq=precond_freq,
-            merge_small=merge_small,
-            max_dim=max_dim,
-            precondition_1d=precondition_1d,
-            eps=eps,
-            decay=decay,
-            bias_correction=bias_correction,
-            alpha=alpha,
-            scale_by_s=scale_by_s,
-            y_to_ema2=y_to_ema2,
-        )
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        updates = []
-        # update preconditioners
-        for i,(p,t, state, setting) in enumerate(zip(params, tensors, states, settings)):
-            beta1, beta2, shampoo_beta, merge_small, max_dim, precondition_1d, eps, alpha = itemgetter(
-                'beta1', 'beta2', 'shampoo_beta', 'merge_small', 'max_dim', 'precondition_1d', 'eps', 'alpha')(setting)
-            scale_by_s = setting['scale_by_s']
-            y_to_ema2 = setting['y_to_ema2']
-
-            if merge_small:
-                t, state['flat_sizes'], state['sort_idxs'] = _merge_small_dims(t, max_dim)
-
-            if 'g_prev' not in state:
-                state['p_prev'] = p.clone()
-                state['g_prev'] = t.clone()
-                updates.append(tensors[i].clip(-0.1,0.1))
-                continue
-
-            p_prev = state['p_prev']
-            g_prev = state['g_prev']
-            s = p - p_prev
-            y = t - g_prev
-            if scale_by_s: y /= torch.linalg.norm(s).clip(min=1e-8) # pylint:disable=not-callable
-
-            state['p_prev'].copy_(p)
-            state['g_prev'].copy_(t)
-
-            # initialize state on 1st step
-            if 'GG' not in state:
-                state["exp_avg"] = torch.zeros_like(t)
-                if y_to_ema2: state["exp_avg_sq"] = torch.ones_like(t)
-                else: state["exp_avg_sq"] = torch.zeros_like(t)
-
-                if not precondition_1d and t.ndim <= 1:
-                    state['GG'] = []
-
-                else:
-                    state['GG'] = [torch.zeros(sh, sh, dtype=t.dtype, device=t.device) if 1<sh<max_dim else None for sh in t.shape]
-
-                # either scalar parameter, 1d with precondition_1d=False, or all dims are too big.
-                if len([i is not None for i in state['GG']]) == 0:
-                    state['GG'] = None
-
-                if state['GG'] is not None:
-                    update_soap_covariances_(y, GGs_=state['GG'], beta=shampoo_beta)
-                    state['Q'] = get_orthogonal_matrix(state['GG'])
-
-                state['step'] = 0
-                updates.append(tensors[i].clip(-0.1,0.1))
-                continue  # skip 1st step as in https://github.com/nikhilvyas/SOAP/blob/main/soap.py ?
-                # I use sign instead as to not mess up with next modules. 1st Adam step is always sign anyway.
-
-            # Projecting gradients to the eigenbases of Shampoo's preconditioner
-            # i.e. projecting to the eigenbases of matrices in state['GG']
-            z_projected = None
-            if state['GG'] is not None:
-                if y_to_ema2: z_projected = project(y, state['Q'])
-                else: z_projected = project(t, state['Q'])
-
-            # exponential moving averages
-            # this part could be foreached but I will do that at some point its not a big difference compared to preconditioning
-            exp_avg: torch.Tensor = state["exp_avg"]
-            exp_avg_sq: torch.Tensor = state["exp_avg_sq"]
-
-            exp_avg.lerp_(t, 1-beta1)
-
-            if z_projected is None:
-                if y_to_ema2: exp_avg_sq.mul_(beta2).addcmul_(y, y, value=1-beta2)
-                else: exp_avg_sq.mul_(beta2).addcmul_(t, t, value=1-beta2)
-            else:
-                exp_avg_sq.mul_(beta2).addcmul_(z_projected, z_projected, value=1-beta2)
-
-            # project exponential moving averages if they are accumulated unprojected
-            exp_avg_projected = exp_avg
-            if z_projected is not None:
-                exp_avg_projected = project(exp_avg, state['Q'])
-
-            exp_avg_sq_projected = exp_avg_sq
-
-            denom = exp_avg_sq_projected.sqrt().add_(eps)
-            # print(f'{t_projected = }, {exp_avg = }, {exp_avg_projected = }, {exp_avg_sq = }, {exp_avg_sq_projected = }, {denom = }')
-
-            # Projecting back the preconditioned (by Adam) exponential moving average of gradients
-            # to the original space
-            update = exp_avg_projected / denom
-            if z_projected is not None:
-                update = project_back(update, state["Q"])
-
-            if setting['bias_correction']:
-                bias_correction1 = 1.0 - beta1 ** (state["step"]+1)
-                bias_correction2 = 1.0 - beta2 ** (state["step"]+1)
-                update *= ((bias_correction2 ** .5) / bias_correction1) * alpha
-            elif alpha is not None:
-                update *= alpha
-
-            if merge_small:
-                update = _unmerge_small_dims(update, state['flat_sizes'], state['sort_idxs'])
-
-            updates.append(update)
-            state["step"] += 1
-
-            # Update is done after the gradient step to avoid using current gradients in the projection.
-            if state['GG'] is not None:
-                update_soap_covariances_(y, state['GG'], shampoo_beta)
-                if state['step'] % setting['precond_freq'] == 0:
-                    state['Q'], state['exp_avg_sq'] = get_orthogonal_matrix_QR(exp_avg_sq, state['GG'], state['Q'])
-
-        return updates

torchzero/modules/experimental/structured_newton.py

@@ -1,111 +0,0 @@
-# idea https://arxiv.org/pdf/2212.09841
-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 (
-    hessian_list_to_mat,
-    hessian_mat,
-    hvp,
-    hvp_fd_central,
-    hvp_fd_forward,
-    jacobian_and_hessian_wrt,
-)
-
-
-class StructuredNewton(Module):
-    """TODO. Please note that this is experimental and isn't guaranteed to work.
-    Args:
-        structure (str, optional): structure.
-        reg (float, optional): tikhonov regularizer value. Defaults to 1e-6.
-        hvp_method (str):
-            how to calculate hvp_method. Defaults to "autograd".
-        inner (Chainable | None, optional): inner modules. Defaults to None.
-
-    """
-    def __init__(
-        self,
-        structure: Literal[
-            "diagonal",
-            "diagonal1",
-            "diagonal_abs",
-            "tridiagonal",
-            "circulant",
-            "toeplitz",
-            "toeplitz_like",
-            "hankel",
-            "rank1",
-            "rank2", # any rank
-        ]
-        | str = "diagonal",
-        reg: float = 1e-6,
-        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
-        h: float = 1e-3,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(reg=reg, hvp_method=hvp_method, structure=structure, h=h)
-        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]]
-        reg = settings['reg']
-        hvp_method = settings['hvp_method']
-        structure = settings['structure']
-        h = settings['h']
-
-        # ------------------------ calculate grad and hessian ------------------------ #
-        if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
-            def Hvp_fn1(x):
-                return hvp(params, grad, x, retain_graph=True)
-            Hvp_fn = Hvp_fn1
-
-        elif hvp_method == 'forward':
-            grad = var.get_grad()
-            def Hvp_fn2(x):
-                return hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1]
-            Hvp_fn = Hvp_fn2
-
-        elif hvp_method == 'central':
-            grad = var.get_grad()
-            def Hvp_fn3(x):
-                return hvp_fd_central(closure, params, x, h=h, normalize=True)[1]
-            Hvp_fn = Hvp_fn3
-
-        else: raise ValueError(hvp_method)
-
-        # -------------------------------- inner step -------------------------------- #
-        update = var.get_update()
-        if 'inner' in self.children:
-            update = apply_transform(self.children['inner'], update, params=params, grads=grad, var=var)
-
-        # hessian
-        if structure.startswith('diagonal'):
-            H = Hvp_fn([torch.ones_like(p) for p in params])
-            if structure == 'diagonal1': torch._foreach_clamp_min_(H, 1)
-            if structure == 'diagonal_abs': torch._foreach_abs_(H)
-            torch._foreach_add_(H, reg)
-            torch._foreach_div_(update, H)
-            var.update = update
-            return var
-
-        # hessian
-        raise NotImplementedError(structure)
-
-
-
-
-

torchzero/modules/lr/__init__.py

@@ -1,2 +0,0 @@
-from .lr import LR, StepSize, Warmup
-from .adaptive import PolyakStepSize, RandomStepSize

torchzero/modules/lr/adaptive.py

@@ -1,93 +0,0 @@
-"""Various step size strategies"""
-import random
-from typing import Any
-from operator import itemgetter
-import torch
-
-from ...core import Transform
-from ...utils import TensorList, NumberList, unpack_dicts
-
-
-class PolyakStepSize(Transform):
-    """Polyak's step-size method.
-
-    Args:
-        max (float | None, optional): maximum possible step size. Defaults to None.
-        min_obj_value (int, optional):
-            (estimated) minimal possible value of the objective function (lowest possible loss). Defaults to 0.
-        use_grad (bool, optional):
-            if True, uses dot product of update and gradient to compute the step size.
-            Otherwise, dot product of update with itself is used, which has no geometric meaning so it probably won't work well.
-            Defaults to True.
-        parameterwise (bool, optional):
-            if True, calculate Polyak step-size for each parameter separately,
-            if False calculate one global step size for all parameters. Defaults to False.
-        alpha (float, optional): multiplier to Polyak step-size. Defaults to 1.
-    """
-    def __init__(self, max: float | None = None, min_obj_value: float = 0, use_grad=True, parameterwise=False, alpha: float = 1):
-
-        defaults = dict(alpha=alpha, max=max, min_obj_value=min_obj_value, use_grad=use_grad, parameterwise=parameterwise)
-        super().__init__(defaults, uses_grad=use_grad)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        assert grads is not None
-        tensors = TensorList(tensors)
-        grads = TensorList(grads)
-        alpha = NumberList(s['alpha'] for s in settings)
-
-        parameterwise, use_grad, max, min_obj_value = itemgetter('parameterwise', 'use_grad', 'max', 'min_obj_value')(settings[0])
-
-        if use_grad: denom = tensors.dot(grads)
-        else: denom = tensors.dot(tensors)
-
-        if parameterwise:
-            polyak_step_size: TensorList | Any = (loss - min_obj_value) / denom.where(denom!=0, 1)
-            polyak_step_size = polyak_step_size.where(denom != 0, 0)
-            if max is not None: polyak_step_size = polyak_step_size.clamp_max(max)
-
-        else:
-            if denom.abs() <= torch.finfo(denom.dtype).eps: polyak_step_size = 0 # converged
-            else: polyak_step_size = (loss - min_obj_value) / denom
-
-            if max is not None:
-                if polyak_step_size > max: polyak_step_size = max
-
-        tensors.mul_(alpha * polyak_step_size)
-        return tensors
-
-
-class RandomStepSize(Transform):
-    """Uses random global or layer-wise step size from `low` to `high`.
-
-    Args:
-        low (float, optional): minimum learning rate. Defaults to 0.
-        high (float, optional): maximum learning rate. Defaults to 1.
-        parameterwise (bool, optional):
-            if True, generate random step size for each parameter separately,
-            if False generate one global random step size. Defaults to False.
-    """
-    def __init__(self, low: float = 0, high: float = 1, parameterwise=False, seed:int|None=None):
-        defaults = dict(low=low, high=high, parameterwise=parameterwise,seed=seed)
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        s = settings[0]
-        parameterwise = s['parameterwise']
-
-        seed = s['seed']
-        if 'generator' not in self.global_state:
-            self.global_state['generator'] = random.Random(seed)
-        generator: random.Random = self.global_state['generator']
-
-        if parameterwise:
-            low, high = unpack_dicts(settings, 'low', 'high')
-            lr = [generator.uniform(l, h) for l, h in zip(low, high)]
-        else:
-            low = s['low']
-            high = s['high']
-            lr = generator.uniform(low, high)
-
-        torch._foreach_mul_(tensors, lr)
-        return tensors

torchzero/modules/lr/lr.py

@@ -1,63 +0,0 @@
-"""Learning rate"""
-import torch
-
-from ...core import Transform
-from ...utils import NumberList, TensorList, generic_eq, unpack_dicts
-
-def lazy_lr(tensors: TensorList, lr: float | list, inplace:bool):
-    """multiplies by lr if lr is not 1"""
-    if generic_eq(lr, 1): return tensors
-    if inplace: return tensors.mul_(lr)
-    return tensors * lr
-
-class LR(Transform):
-    """Learning rate. Adding this module also adds support for LR schedulers."""
-    def __init__(self, lr: float):
-        defaults=dict(lr=lr)
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        return lazy_lr(TensorList(tensors), lr=[s['lr'] for s in settings], inplace=True)
-
-class StepSize(Transform):
-    """this is exactly the same as LR, except the `lr` parameter can be renamed to any other name to avoid clashes"""
-    def __init__(self, step_size: float, key = 'step_size'):
-        defaults={"key": key, key: step_size}
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        return lazy_lr(TensorList(tensors), lr=[s[s['key']] for s in settings], inplace=True)
-
-
-def _warmup_lr(step: int, start_lr: float | NumberList, end_lr: float | NumberList, steps: float):
-    """returns warm up lr scalar"""
-    if step > steps: return end_lr
-    return start_lr + (end_lr - start_lr) * (step / steps)
-
-class Warmup(Transform):
-    """Learning rate warmup, linearly increases learning rate multiplier from :code:`start_lr` to :code:`end_lr` over :code:`steps` steps.
-
-    Args:
-        start_lr (_type_, optional): initial learning rate multiplier on first step. Defaults to 1e-5.
-        end_lr (float, optional): learning rate multiplier at the end and after warmup. Defaults to 1.
-        steps (int, optional): number of steps to perform warmup for. Defaults to 100.
-    """
-    def __init__(self, start_lr = 1e-5, end_lr:float = 1, steps = 100):
-        defaults = dict(start_lr=start_lr,end_lr=end_lr, steps=steps)
-        super().__init__(defaults, uses_grad=False)
-
-    @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
-        start_lr, end_lr = unpack_dicts(settings, 'start_lr', 'end_lr', cls = NumberList)
-        num_steps = settings[0]['steps']
-        step = self.global_state.get('step', 0)
-
-        target = lazy_lr(
-            TensorList(tensors),
-            lr=_warmup_lr(step=step, start_lr=start_lr, end_lr=end_lr, steps=num_steps),
-            inplace=True
-        )
-        self.global_state['step'] = step + 1
-        return target