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

torchzero/modules/experimental/eigendescent.py

@@ -1,120 +0,0 @@
-from contextlib import nullcontext
-import warnings
-from collections.abc import Callable
-from functools import partial
-import itertools
-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,
-    jacobian_wrt, jacobian_and_hessian_wrt, hessian_mat,
-)
-
-def _batched_dot(x, y):
-    return (x.unsqueeze(-2) @ y.unsqueeze(-1)).squeeze(-1).squeeze(-1)
-
-def _cosine_similarity(x, y):
-    denom = torch.linalg.vector_norm(x, dim=-1) * torch.linalg.vector_norm(y, dim=-1).clip(min=torch.finfo(x.dtype).eps) # pylint:disable=not-callable
-    return _batched_dot(x, y) / denom
-
-class EigenDescent(Module):
-    """
-    Uses eigenvectors corresponding to certain eigenvalues. For now they are just extracted from hessian.
-
-    .. warning::
-        Experimental.
-
-    Args:
-        mode (str, optional):
-            - largest - use largest eigenvalue unless all eigenvalues are negative, then smallest is used.
-            - smallest - use smallest eigenvalue unless all eigenvalues are positive, then largest is used.
-            - mean-sign - use mean of eigenvectors multiplied by 1 or -1 if they point in opposite direction from gradient.
-            - mean-dot - use mean of eigenvectors multiplied by dot product with gradient.
-            - mean-cosine - use mean of eigenvectors multiplied by cosine similarity with gradient.
-            - mm - for testing.
-
-            Defaults to 'mean-sign'.
-        hessian_method (str, optional): how to calculate hessian. Defaults to "autograd".
-        vectorize (bool, optional): how to calculate hessian. Defaults to True.
-
-    """
-    def __init__(
-        self,
-        mode: Literal['largest', 'smallest','magnitude', 'mean-sign', 'mean-dot', 'mean-cosine', 'mm'] = 'mean-sign',
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
-    ):
-        defaults = dict(hessian_method=hessian_method, vectorize=vectorize, mode=mode)
-        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]]
-        mode = settings['mode']
-        hessian_method = settings['hessian_method']
-        vectorize = settings['vectorize']
-
-        # ------------------------ calculate grad and hessian ------------------------ #
-        if hessian_method == 'autograd':
-            with torch.enable_grad():
-                loss = var.loss = var.loss_approx = closure(False)
-                g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
-                g_list = [t[0] for t in g_list] # remove leading dim from loss
-                var.grad = g_list
-                H = hessian_list_to_mat(H_list)
-
-        elif hessian_method in ('func', 'autograd.functional'):
-            strat = 'forward-mode' if vectorize else 'reverse-mode'
-            with torch.enable_grad():
-                g_list = var.get_grad(retain_graph=True)
-                H: torch.Tensor = hessian_mat(partial(closure, backward=False), params,
-                                method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
-
-        else:
-            raise ValueError(hessian_method)
-
-
-        # ----------------------------------- solve ---------------------------------- #
-        g = torch.cat([t.ravel() for t in g_list])
-        L, Q = torch.linalg.eigh(H) # L is sorted # pylint:disable=not-callable
-        if mode == 'largest':
-            # smallest eigenvalue if all eigenvalues are negative else largest
-            if L[-1] <= 0: d = Q[0]
-            else: d = Q[-1]
-
-        elif mode == 'smallest':
-            # smallest eigenvalue if negative eigenvalues exist else largest
-            if L[0] <= 0: d = Q[0]
-            else: d = Q[-1]
-
-        elif mode == 'magnitude':
-            # largest by magnitude
-            if L[0].abs() > L[-1].abs(): d = Q[0]
-            else: d = Q[-1]
-
-        elif mode == 'mean-dot':
-            d = ((g.unsqueeze(0) @ Q).squeeze(0) * Q).mean(1)
-
-        elif mode == 'mean-sign':
-            d = ((g.unsqueeze(0) @ Q).squeeze(0).sign() * Q).mean(1)
-
-        elif mode == 'mean-cosine':
-            d = (Q * _cosine_similarity(Q, g)).mean(1)
-
-        elif mode == 'mm':
-            d = (g.unsqueeze(0) @ Q).squeeze(0) / g.numel()
-
-        else:
-            raise ValueError(mode)
-
-        var.update = vec_to_tensors(g.dot(d).sign() * d, params)
-        return var
-

torchzero/modules/experimental/etf.py

@@ -1,195 +0,0 @@
-from typing import cast
-import warnings
-
-import torch
-
-from ...core import Module
-from ...utils import vec_to_tensors, vec_to_tensors_, as_tensorlist
-
-
-class ExponentialTrajectoryFit(Module):
-    """A method.
-
-    .. warning::
-        Experimental.
-    """
-    def __init__(self, step_size=1e-2, adaptive:bool=True):
-        defaults = dict(step_size = step_size,adaptive=adaptive)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        assert closure is not None
-        step_size = self.settings[var.params[0]]['step_size']
-        adaptive = self.settings[var.params[0]]['adaptive']
-
-
-        # 1. perform 3 GD steps to obtain 4 points
-        points = [torch.cat([p.view(-1) for p in var.params])]
-        for i in range(3):
-            if i == 0:
-                grad = var.get_grad()
-                if adaptive:
-                    step_size /= as_tensorlist(grad).abs().global_mean().clip(min=1e-4)
-
-            else:
-                with torch.enable_grad(): closure()
-                grad = [cast(torch.Tensor, p.grad) for p in var.params]
-
-            # GD step
-            torch._foreach_sub_(var.params, grad, alpha=step_size)
-
-            points.append(torch.cat([p.view(-1) for p in var.params]))
-
-        assert len(points) == 4, len(points)
-        x0, x1, x2, x3 = points
-        dim = x0.numel()
-
-        # 2. fit a generalized exponential curve
-        d0 = (x1 - x0).unsqueeze(1) # column vectors
-        d1 = (x2 - x1).unsqueeze(1)
-        d2 = (x3 - x2).unsqueeze(1)
-
-        # cat
-        D1 = torch.cat([d0, d1], dim=1)
-        D2 = torch.cat([d1, d2], dim=1)
-
-        # if points are collinear this will happen on sphere and a quadratic "line search" will minimize it
-        if x0.numel() >= 2:
-            if torch.linalg.matrix_rank(D1) < 2: # pylint:disable=not-callable
-                pass # need to put a quadratic fit there
-
-        M = D2 @ torch.linalg.pinv(D1) # pylint:disable=not-callable # this defines the curve
-
-        # now we can predict x*
-        I = torch.eye(dim, device=x0.device, dtype=x0.dtype)
-        B = I - M
-        z = x1 - M @ x0
-
-        x_star = torch.linalg.lstsq(B, z).solution # pylint:disable=not-callable
-
-        vec_to_tensors_(x0, var.params)
-        difference = torch._foreach_sub(var.params, vec_to_tensors(x_star, var.params))
-        var.update = list(difference)
-        return var
-
-
-
-class ExponentialTrajectoryFitV2(Module):
-    """Should be better than one above, except it isn't.
-
-    .. warning::
-        Experimental.
-
-    """
-    def __init__(self, step_size=1e-3, num_steps: int= 4, adaptive:bool=True):
-        defaults = dict(step_size = step_size, num_steps=num_steps, adaptive=adaptive)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        assert closure is not None
-        step_size = self.settings[var.params[0]]['step_size']
-        num_steps = self.settings[var.params[0]]['num_steps']
-        adaptive = self.settings[var.params[0]]['adaptive']
-
-        # 1. perform 3 GD steps to obtain 4 points (or more)
-        grad = var.get_grad()
-        if adaptive:
-            step_size /= as_tensorlist(grad).abs().global_mean().clip(min=1e-4)
-
-        points = [torch.cat([p.view(-1) for p in var.params])]
-        point_grads = [torch.cat([g.view(-1) for g in grad])]
-
-        for i in range(num_steps):
-            # GD step
-            torch._foreach_sub_(var.params, grad, alpha=step_size)
-
-            points.append(torch.cat([p.view(-1) for p in var.params]))
-
-            closure(backward=True)
-            grad = [cast(torch.Tensor, p.grad) for p in var.params]
-            point_grads.append(torch.cat([g.view(-1) for g in grad]))
-
-
-        X = torch.stack(points, 1) # dim, num_steps+1
-        G = torch.stack(point_grads, 1)
-        dim = points[0].numel()
-
-        X = torch.cat([X, torch.ones(1, num_steps+1, dtype=G.dtype, device=G.device)])
-
-        P = G @ torch.linalg.pinv(X) # pylint:disable=not-callable
-        A = P[:, :dim]
-        b = -P[:, dim]
-
-        # symmetrize
-        A = 0.5 * (A + A.T)
-
-        # predict x*
-        x_star = torch.linalg.lstsq(A, b).solution # pylint:disable=not-callable
-
-        vec_to_tensors_(points[0], var.params)
-        difference = torch._foreach_sub(var.params, vec_to_tensors(x_star, var.params))
-        var.update = list(difference)
-        return var
-
-
-
-
-def _fit_exponential(y0, y1, y2):
-    """x0, x1 and x2 are assumed to be 0, 1, 2"""
-    r = (y2 - y1) / (y1 - y0)
-    ones = r==1
-    r[ones] = 0
-    B = (y1 - y0) / (r - 1)
-    A = y0 - B
-
-    A[ones] = 0
-    B[ones] = 0
-    return A, B, r
-
-class PointwiseExponential(Module):
-    """A stupid method (for my youtube channel).
-
-    .. warning::
-        Experimental.
-    """
-    def __init__(self, step_size: float = 1e-3, reg: float = 1e-2, steps = 10000):
-        defaults = dict(reg=reg, steps=steps, step_size=step_size)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        assert closure is not None
-        settings = self.settings[var.params[0]]
-        step_size = settings['step_size']
-        reg = settings['reg']
-        steps = settings['steps']
-
-        # 1. perform 2 GD steps to obtain 3 points
-        points = [torch.cat([p.view(-1) for p in var.params])]
-        for i in range(2):
-            if i == 0: grad = var.get_grad()
-            else:
-                with torch.enable_grad(): closure()
-                grad = [cast(torch.Tensor, p.grad) for p in var.params]
-
-            # GD step
-            torch._foreach_sub_(var.params, grad, alpha=step_size)
-
-            points.append(torch.cat([p.view(-1) for p in var.params]))
-
-        assert len(points) == 3, len(points)
-        y0, y1, y2 = points
-
-        A, B, r = _fit_exponential(y0, y1, y2)
-        r = r.clip(max = 1-reg)
-        x_star = A + B * r**steps
-
-        vec_to_tensors_(y0, var.params)
-        difference = torch._foreach_sub(var.params, vec_to_tensors(x_star, var.params))
-        var.update = list(difference)
-        return var

torchzero/modules/experimental/exp_adam.py

@@ -1,113 +0,0 @@
-from operator import itemgetter
-from functools import partial
-import math
-import torch
-
-from ...core import Module, Target, Transform, apply_transform, Chainable
-from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
-from ..functional import (
-    debias, debiased_step_size,
-    ema_,
-    sqrt_ema_sq_,
-)
-from ..step_size.lr import lazy_lr
-from ..momentum.experimental import sqrt_nag_ema_sq_
-from ..momentum.momentum import nag_
-
-
-def exp_adam_(
-    tensors: TensorList,
-    exp_avg_: TensorList,
-    exp_avg_exp_: TensorList,
-    alpha: float | NumberList,
-    beta1: float | NumberList,
-    beta2: float | NumberList,
-    eps: float | NumberList,
-    step: int,
-    pow: float = 2,
-    debiased: bool = True,
-    max_exp_avg_exp_: TensorList | None = None,
-
-    # inner args
-    inner: Module | None = None,
-    params: list[torch.Tensor] | None = None,
-    grads: list[torch.Tensor] | None = None,
-):
-    """Returns new tensors."""
-    tensors_exp = tensors.abs().clip_(max=math.log(torch.finfo(tensors[0].dtype).max) / 2).exp_()
-    exp_avg_exp_.lerp_(tensors_exp, 1-beta2)
-
-    if max_exp_avg_exp_ is not None:
-        max_exp_avg_exp_.maximum_(exp_avg_exp_)
-        exp_avg_exp_ = max_exp_avg_exp_
-
-    if inner is not None:
-        assert params is not None
-        tensors = TensorList(apply_transform(inner, tensors, params=params, grads=grads))
-
-    exp_avg_ = ema_(tensors, exp_avg_=exp_avg_, beta=beta1, dampening=0,lerp=True)
-    if debiased: alpha = debiased_step_size(step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
-    return (exp_avg_.lazy_mul(alpha) / exp_avg_exp_.log().add_(eps))
-
-class ExpAdam(Transform):
-    """Adam but uses abs exp and log instead of square and sqrt.
-    The gradient will be clipped to half the maximum value representable by its dtype (around 50 for float32)
-
-    Args:
-        beta1 (float, optional): momentum. Defaults to 0.9.
-        beta2 (float, optional): second momentum. Defaults to 0.999.
-        eps (float, optional): epsilon. Defaults to 1e-8.
-        alpha (float, optional): learning rate. Defaults to 1.
-        amsgrad (bool, optional): Whether to divide by maximum of EMA of gradient squares instead. Defaults to False.
-        pow (float, optional): power used in second momentum power and root. Defaults to 2.
-        debiased (bool, optional): whether to apply debiasing to momentums based on current step. Defaults to True.
-    """
-    def __init__(
-        self,
-        beta1: float = 0.9,
-        beta2: float = 0.999,
-        eps: float = 1e-8,
-        amsgrad: bool = False,
-        alpha: float = 1.,
-        pow: float = 2,
-        debiased: bool = True,
-        inner: Chainable | None = None
-    ):
-        defaults=dict(beta1=beta1,beta2=beta2,eps=eps,alpha=alpha,amsgrad=amsgrad,pow=pow,debiased=debiased)
-        super().__init__(defaults, uses_grad=False)
-
-        if inner is not None: self.set_child('inner', inner)
-
-    @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
-        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
-
-        beta1,beta2,eps,alpha=unpack_dicts(settings, 'beta1','beta2','eps','alpha', cls=NumberList)
-        amsgrad,pow,debiased = itemgetter('amsgrad','pow','debiased')(settings[0])
-
-        if amsgrad:
-            exp_avg, exp_avg_exp, max_exp_avg_exp = unpack_states(states, tensors, 'exp_avg', 'exp_avg_exp', 'max_exp_avg_exp', cls=TensorList)
-        else:
-            exp_avg, exp_avg_exp = unpack_states(states, tensors, 'exp_avg', 'exp_avg_exp', cls=TensorList)
-            max_exp_avg_exp = None
-
-
-        return exp_adam_(
-            tensors=TensorList(tensors),
-            exp_avg_=exp_avg,
-            exp_avg_exp_=exp_avg_exp,
-            alpha=alpha,
-            beta1=beta1,
-            beta2=beta2,
-            eps=eps,
-            step=step,
-            pow=pow,
-            debiased=debiased,
-            max_exp_avg_exp_=max_exp_avg_exp,
-
-            # inner args
-            inner=self.children.get("inner", None),
-            params=params,
-            grads=grads,
-
-        )

torchzero/modules/experimental/expanded_lbfgs.py

@@ -1,141 +0,0 @@
-from collections import deque
-from operator import itemgetter
-import torch
-
-from ...core import Transform, Chainable, Module, Var, apply_transform
-from ...utils import TensorList, as_tensorlist, NumberList
-from ...modules.quasi_newton.lbfgs import _adaptive_damping, lbfgs, _lerp_params_update_
-
-class ExpandedLBFGS(Module):
-    """L-BFGS but uses differences between more pairs than just consequtive. Window size controls how far away the pairs are allowed to be.
-    """
-    def __init__(
-        self,
-        history_size=10,
-        window_size:int=3,
-        tol: float | None = 1e-10,
-        damping: bool = False,
-        init_damping=0.9,
-        eigval_bounds=(0.5, 50),
-        params_beta: float | None = None,
-        grads_beta: float | None = None,
-        update_freq = 1,
-        z_beta: float | None = None,
-        tol_reset: bool = False,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(history_size=history_size, window_size=window_size, tol=tol, damping=damping, init_damping=init_damping, eigval_bounds=eigval_bounds, params_beta=params_beta, grads_beta=grads_beta, update_freq=update_freq, z_beta=z_beta, tol_reset=tol_reset)
-        super().__init__(defaults)
-
-        self.global_state['s_history'] = deque(maxlen=history_size)
-        self.global_state['y_history'] = deque(maxlen=history_size)
-        self.global_state['sy_history'] = deque(maxlen=history_size)
-        self.global_state['p_history'] = deque(maxlen=window_size)
-        self.global_state['g_history'] = deque(maxlen=window_size)
-
-        if inner is not None:
-            self.set_child('inner', inner)
-
-    def reset(self):
-        self.state.clear()
-        self.global_state['step'] = 0
-        self.global_state['s_history'].clear()
-        self.global_state['y_history'].clear()
-        self.global_state['sy_history'].clear()
-        self.global_state['p_history'].clear()
-        self.global_state['g_history'].clear()
-
-    @torch.no_grad
-    def step(self, var):
-        params = as_tensorlist(var.params)
-        update = as_tensorlist(var.get_update())
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
-
-        # history of s and k
-        s_history: deque[TensorList] = self.global_state['s_history']
-        y_history: deque[TensorList] = self.global_state['y_history']
-        sy_history: deque[torch.Tensor] = self.global_state['sy_history']
-        p_history: deque[TensorList] = self.global_state['p_history']
-        g_history: deque[TensorList] = self.global_state['g_history']
-
-        tol, damping, init_damping, eigval_bounds, update_freq, z_beta, tol_reset = itemgetter(
-            'tol', 'damping', 'init_damping', 'eigval_bounds', 'update_freq', 'z_beta', 'tol_reset')(self.settings[params[0]])
-        params_beta, grads_beta = self.get_settings(params, 'params_beta', 'grads_beta')
-
-        l_params, l_update = _lerp_params_update_(self, params, update, params_beta, grads_beta)
-        prev_l_params, prev_l_grad = self.get_state(params, 'prev_l_params', 'prev_l_grad', cls=TensorList)
-
-        # 1st step - there are no previous params and grads, lbfgs will do normalized GD step
-        if step == 0:
-            s = None; y = None; ys = None
-        else:
-            s = l_params - prev_l_params
-            y = l_update - prev_l_grad
-            ys = s.dot(y)
-
-            if damping:
-                s, y, ys = _adaptive_damping(s, y, ys, init_damping=init_damping, eigval_bounds=eigval_bounds)
-
-        prev_l_params.copy_(l_params)
-        prev_l_grad.copy_(l_update)
-
-        # update effective preconditioning state
-        if step % update_freq == 0:
-            if ys is not None and ys > 1e-10:
-                assert s is not None and y is not None
-                s_history.append(s)
-                y_history.append(y)
-                sy_history.append(ys)
-
-            if len(p_history) > 1:
-                for p_i, g_i in zip(list(p_history)[:-1], list(g_history)[:-1]):
-                    s_i = l_params - p_i
-                    y_i = l_update - g_i
-                    ys_i = s_i.dot(y_i)
-
-                    if ys_i > 1e-10:
-                        if damping:
-                            s_i, y_i, ys_i = _adaptive_damping(s_i, y_i, ys_i, init_damping=init_damping, eigval_bounds=eigval_bounds)
-
-                        s_history.append(s_i)
-                        y_history.append(y_i)
-                        sy_history.append(ys_i)
-
-            p_history.append(l_params.clone())
-            g_history.append(l_update.clone())
-
-
-        # step with inner module before applying preconditioner
-        if self.children:
-            update = TensorList(apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var))
-
-        # tolerance on gradient difference to avoid exploding after converging
-        if tol is not None:
-            if y is not None and y.abs().global_max() <= tol:
-                var.update = update # may have been updated by inner module, probably makes sense to use it here?
-                if tol_reset: self.reset()
-                return var
-
-        # lerp initial H^-1 @ q guess
-        z_ema = None
-        if z_beta is not None:
-            z_ema = self.get_state(var.params, 'z_ema', cls=TensorList)
-
-        # precondition
-        dir = lbfgs(
-            tensors_=as_tensorlist(update),
-            s_history=s_history,
-            y_history=y_history,
-            sy_history=sy_history,
-            y=y,
-            sy=ys,
-            z_beta = z_beta,
-            z_ema = z_ema,
-            step=step
-        )
-
-        var.update = dir
-
-        return var
-

torchzero/modules/experimental/hnewton.py

@@ -1,85 +0,0 @@
-from collections import deque
-
-import torch
-
-from ...core import TensorwiseTransform
-
-
-def eigh_solve(H: torch.Tensor, g: torch.Tensor):
-    try:
-        L, Q = torch.linalg.eigh(H) # pylint:disable=not-callable
-        return Q @ ((Q.mH @ g) / L)
-    except torch.linalg.LinAlgError:
-        return None
-
-
-class HNewton(TensorwiseTransform):
-    """This treats gradient differences as Hvps with vectors being parameter differences, using past gradients that are close to each other. Basically this is another limited memory quasi newton method to test.
-
-    .. warning::
-        Experimental.
-
-    """
-    def __init__(self, history_size: int, window_size: int, reg: float=0, tol: float = 1e-8, concat_params:bool=True, inner=None):
-        defaults = dict(history_size=history_size, window_size=window_size, reg=reg, tol=tol)
-        super().__init__(defaults, uses_grad=False, concat_params=concat_params, inner=inner)
-
-    def update_tensor(self, tensor, param, grad, loss, state, setting):
-
-        history_size = setting['history_size']
-
-        if 'param_history' not in state:
-            state['param_history'] = deque(maxlen=history_size)
-            state['grad_history'] = deque(maxlen=history_size)
-
-        param_history: deque = state['param_history']
-        grad_history: deque = state['grad_history']
-        param_history.append(param.ravel())
-        grad_history.append(tensor.ravel())
-
-    def apply_tensor(self, tensor, param, grad, loss, state, setting):
-        window_size = setting['window_size']
-        reg = setting['reg']
-        tol = setting['tol']
-
-        param_history: deque = state['param_history']
-        grad_history: deque = state['grad_history']
-        g = tensor.ravel()
-
-        n = len(param_history)
-        s_list = []
-        y_list = []
-
-        for i in range(n):
-            for j in range(i):
-                if i - j <= window_size:
-                    p_i, g_i = param_history[i], grad_history[i]
-                    p_j, g_j = param_history[j], grad_history[j]
-                    s = p_i - p_j # vec in hvp
-                    y = g_i - g_j # hvp
-                    if s.dot(y) > tol:
-                        s_list.append(s)
-                        y_list.append(y)
-
-        if len(s_list) < 1:
-            scale = (1 / tensor.abs().sum()).clip(min=torch.finfo(tensor.dtype).eps, max=1)
-            tensor.mul_(scale)
-            return tensor
-
-        S = torch.stack(s_list, 1)
-        Y = torch.stack(y_list, 1)
-
-        B = S.T @ Y
-        if reg != 0: B.add_(torch.eye(B.size(0), device=B.device, dtype=B.dtype).mul_(reg))
-        g_proj = g @ S
-
-        newton_proj, info = torch.linalg.solve_ex(B, g_proj) # pylint:disable=not-callable
-        if info != 0:
-            newton_proj = -torch.linalg.lstsq(B, g_proj).solution # pylint:disable=not-callable
-        newton = S @ newton_proj
-        return newton.view_as(tensor)
-
-
-        # scale = (1 / tensor.abs().sum()).clip(min=torch.finfo(tensor.dtype).eps, max=1)
-        # tensor.mul_(scale)
-        # return tensor