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

torchzero/modules/{momentum/experimental.py → experimental/momentum.py}

@@ -6,10 +6,10 @@ from typing import Literal
 import torch
 
 from ...core import Target, Transform
-from ...utils import NumberList, TensorList, unpack_states, unpack_dicts
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 from ..functional import ema_, ema_sq_, sqrt_ema_sq_
-from .ema import EMASquared, SqrtEMASquared
-from .momentum import nag_
+from ..momentum.momentum import nag_
+from ..ops.higher_level import EMASquared, SqrtEMASquared
 
 
 def precentered_ema_sq_(
@@ -49,7 +49,7 @@ class PrecenteredEMASquared(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
 
         beta1, beta2 = unpack_dicts(settings, 'beta1','beta2', cls=NumberList)
@@ -154,44 +154,7 @@ class CoordinateMomentum(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         p = NumberList(s['p'] for s in settings)
         velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
         return coordinate_momentum_(TensorList(tensors), velocity_=velocity, p=p).clone()
-
-
-# def multiplicative_momentum_(
-#     tensors_: TensorList,
-#     velocity_: TensorList,
-#     momentum: float | NumberList,
-#     dampening: float | NumberList,
-#     normalize_velocity: bool = True,
-#     abs: bool = False,
-#     lerp: bool = False,
-# ):
-#     """
-#     abs: if True, tracks momentum of absolute magnitudes.
-
-#     returns `tensors_`.
-#     """
-#     tensors_into_velocity = tensors_.abs() if abs else tensors_
-#     ema_(tensors_into_velocity, exp_avg_=velocity_, beta=momentum, dampening=0, lerp=lerp)
-
-#     if normalize_velocity: velocity_ = velocity_ / velocity_.std().add_(1e-8)
-#     return tensors_.mul_(velocity_.lazy_mul(1-dampening) if abs else velocity_.abs().lazy_mul_(1-dampening))
-
-
-# class MultiplicativeMomentum(Transform):
-#     """sucks"""
-#     def __init__(self, momentum: float = 0.9, dampening: float = 0,normalize_velocity: bool = True, abs: bool = False, lerp: bool = False):
-#         defaults = dict(momentum=momentum, dampening=dampening, normalize_velocity=normalize_velocity,abs=abs, lerp=lerp)
-#         super().__init__(defaults, uses_grad=False)
-
-#     @torch.no_grad
-#     def apply(self, tensors, params, grads, loss, states, settings):
-#         momentum,dampening = self.get_settings('momentum','dampening', params=params, cls=NumberList)
-#         abs,lerp,normalize_velocity = self.first_setting('abs','lerp','normalize_velocity', params=params)
-#         velocity = self.get_state('velocity', params=params, cls=TensorList)
-#         return multiplicative_momentum_(TensorList(target), velocity_=velocity, momentum=momentum, dampening=dampening,
-#                                         normalize_velocity=normalize_velocity,abs=abs,lerp=lerp)
-

torchzero/modules/experimental/newton_solver.py

@@ -3,28 +3,36 @@ from typing import Any, Literal, overload
 
 import torch
 
-from ...core import Chainable, Module, apply_transform, Modular
+from ...core import Chainable, Modular, Module, apply_transform
 from ...utils import TensorList, as_tensorlist
-from ...utils.derivatives import hvp
+from ...utils.derivatives import hvp, hvp_fd_forward, hvp_fd_central
 from ..quasi_newton import LBFGS
 
+
 class NewtonSolver(Module):
-    """Matrix free newton via with any custom solver (this is for testing, use NewtonCG or NystromPCG)"""
+    """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()),
         maxiter=None,
-        tol=1e-3,
+        maxiter1=None,
+        tol:float | None=1e-3,
         reg: float = 0,
         warm_start=True,
+        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        reset_solver: bool = False,
+        h: float= 1e-3,
         inner: Chainable | None = None,
     ):
-        defaults = dict(tol=tol, maxiter=maxiter, reg=reg, warm_start=warm_start, solver=solver)
+        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,)
 
         if inner is not None:
             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)
@@ -34,19 +42,49 @@ class NewtonSolver(Module):
         settings = self.settings[params[0]]
         solver_cls = settings['solver']
         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 --------------------- #
-        grad = var.get_grad(create_graph=True)
+        if hvp_method == 'autograd':
+            grad = var.get_grad(create_graph=True)
 
-        def H_mm(x):
-            with torch.enable_grad():
-                Hvp = TensorList(hvp(params, grad, x, 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)
+
+
         # -------------------------------- inner step -------------------------------- #
         b = as_tensorlist(grad)
         if 'inner' in self.children:
@@ -58,23 +96,46 @@ class NewtonSolver(Module):
         if x0 is None: x = b.zeros_like().requires_grad_(True)
         else: x = x0.clone().requires_grad_(True)
 
-        solver = solver_cls(x)
+
+        if 'solver' not in self.global_state:
+            if maxiter1 is not None: maxiter = maxiter1
+            solver = self.global_state['solver'] = solver_cls(x)
+            self.global_state['x'] = x
+
+        else:
+            if reset_solver:
+                solver = self.global_state['solver'] = solver_cls(x)
+            else:
+                solver_params = self.global_state['x']
+                solver_params.set_(x)
+                x = solver_params
+                solver = self.global_state['solver']
+
         def lstsq_closure(backward=True):
-            Hx = H_mm(x)
-            loss = (Hx-b).pow(2).global_mean()
+            Hx = H_mm(x).detach()
+            # loss = (Hx-b).pow(2).global_mean()
+            # if backward:
+            #     solver.zero_grad()
+            #     loss.backward(inputs=x)
+
+            residual = Hx - b
+            loss = residual.pow(2).global_mean()
             if backward:
-                solver.zero_grad()
-                loss.backward(inputs=x)
+                with torch.no_grad():
+                    H_residual = H_mm(residual)
+                    n = residual.global_numel()
+                    x.set_grad_((2.0 / n) * H_residual)
+
             return loss
 
         if maxiter is None: maxiter = b.global_numel()
         loss = None
-        initial_loss = lstsq_closure(False)
-        if initial_loss > tol:
+        initial_loss = lstsq_closure(False) if tol is not None else None # skip unnecessary closure if tol is None
+        if initial_loss is None or initial_loss > torch.finfo(b[0].dtype).eps:
             for i in range(maxiter):
                 loss = solver.step(lstsq_closure)
                 assert loss is not None
-                if min(loss, loss/initial_loss) < tol: break
+                if initial_loss is not None and loss/initial_loss < tol: break
 
         # print(f'{loss = }')
 
@@ -83,6 +144,7 @@ class NewtonSolver(Module):
             x0.copy_(x)
 
         var.update = x.detach()
+        self._num_hvps += self._num_hvps_last_step
         return var
 
 

torchzero/modules/experimental/newtonnewton.py

@@ -10,20 +10,21 @@ import torch
 from ...core import Chainable, Module, apply_transform
 from ...utils import TensorList, vec_to_tensors
 from ...utils.derivatives import (
-    hessian_list_to_mat,
+    flatten_jacobian,
     jacobian_wrt,
 )
 from ..second_order.newton import (
-    cholesky_solve,
-    eigh_solve,
-    least_squares_solve,
-    lu_solve,
+    _cholesky_solve,
+    _eigh_solve,
+    _least_squares_solve,
+    _lu_solve,
 )
-
+from ...utils.linalg.linear_operator import Dense
 
 class NewtonNewton(Module):
-    """
-    Method that I thought of and then it worked.
+    """Applies Newton-like preconditioning to Newton step.
+
+    This is a method that I thought of and then it worked. Here is how it works:
 
     1. Calculate newton step by solving Hx=g
 
@@ -34,6 +35,9 @@ class NewtonNewton(Module):
     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,
@@ -47,10 +51,10 @@ class NewtonNewton(Module):
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, var):
+    def update(self, var):
         params = TensorList(var.params)
         closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
+        if closure is None: raise RuntimeError('NewtonNewton requires closure')
 
         settings = self.settings[params[0]]
         reg = settings['reg']
@@ -60,6 +64,7 @@ class NewtonNewton(Module):
         eigval_tfm = settings['eigval_tfm']
 
         # ------------------------ calculate grad and hessian ------------------------ #
+        Hs = []
         with torch.enable_grad():
             loss = var.loss = var.loss_approx = closure(False)
             g_list = torch.autograd.grad(loss, params, create_graph=True)
@@ -72,17 +77,29 @@ class NewtonNewton(Module):
                 is_last = o == order
                 H_list = jacobian_wrt([xp], params, create_graph=not is_last, batched=vectorize)
                 with torch.no_grad() if is_last else nullcontext():
-                    H = hessian_list_to_mat(H_list)
+                    H = flatten_jacobian(H_list)
                     if reg != 0: H = H + I * reg
+                    Hs.append(H)
 
                     x = None
                     if search_negative or (is_last and eigval_tfm is not None):
-                        x = eigh_solve(H, xp, eigval_tfm, 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)
+                        x = _eigh_solve(H, xp, eigval_tfm, 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()
 
+        self.global_state["Hs"] = Hs
+        self.global_state['xp'] = xp.nan_to_num_(0,0,0)
+
+    @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

torchzero/modules/experimental/reduce_outward_lr.py

@@ -4,19 +4,19 @@ from ...core import Target, Transform
 from ...utils import TensorList, unpack_states, unpack_dicts
 
 class ReduceOutwardLR(Transform):
-    """
-    When update sign matches weight sign, the learning rate for that weight is multiplied by `mul`.
+    """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.
 
-    A note on this is that it sounded good but its really bad in practice.
+    .. 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'):
         defaults = dict(mul=mul, use_grad=use_grad, invert=invert)
         super().__init__(defaults, uses_grad=use_grad, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         params = TensorList(params)
         tensors = TensorList(tensors)
 

torchzero/modules/experimental/scipy_newton_cg.py

@@ -0,0 +1,105 @@
+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
+
+
+class ScipyNewtonCG(Module):
+    """NewtonCG with scipy solvers (any from scipy.sparse.linalg)"""
+    def __init__(
+        self,
+        solver = gcrotmk,
+        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        h: float = 1e-3,
+        warm_start=False,
+        inner: Chainable | None = None,
+        kwargs: dict | None = None,
+    ):
+        defaults = dict(hvp_method=hvp_method, solver=solver, h=h, warm_start=warm_start)
+        super().__init__(defaults,)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+        self._num_hvps = 0
+        self._num_hvps_last_step = 0
+
+        if kwargs is None: kwargs = {}
+        self._kwargs = kwargs
+
+    @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]]
+        hvp_method = settings['hvp_method']
+        solver = settings['solver']
+        h = settings['h']
+        warm_start = settings['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)
+
+            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))
+                return torch.cat([t.ravel() for t in Hvp]).numpy(force=True)
+
+        else:
+
+            with torch.enable_grad():
+                grad = var.get_grad()
+
+            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])
+                    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])
+                    return torch.cat([t.ravel() for t in Hvp]).numpy(force=True)
+
+            else:
+                raise ValueError(hvp_method)
+
+        ndim = sum(p.numel() for p in params)
+        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)
+
+        # ---------------------------------- run cg ---------------------------------- #
+        x0 = None
+        if warm_start: x0 = self.global_state.get('x_prev', None) # initialized to 0 which is default anyway
+
+        x_np = solver(H, b.to_vec().nan_to_num().numpy(force=True), x0=x0, **self._kwargs)
+        if isinstance(x_np, tuple): x_np = x_np[0]
+
+        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)
+
+        self._num_hvps += self._num_hvps_last_step
+        return var
+

torchzero/modules/{projections/structural.py → experimental/structural_projections.py}

@@ -5,36 +5,19 @@ import torch
 
 from ...core import Chainable
 from ...utils import vec_to_tensors, TensorList
-from ..optimizers.shampoo import _merge_small_dims
-from .projection import Projection
+from ..adaptive.shampoo import _merge_small_dims
+from ..projections import ProjectionBase
 
 
-class VectorProjection(Projection):
-    """
-    flattens and concatenates all parameters into a vector
-    """
-    def __init__(self, modules: Chainable, project_update=True, project_params=False, project_grad=False):
-        super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
-    @torch.no_grad
-    def project(self, tensors, var, current):
-        return [torch.cat([u.view(-1) for u in tensors], dim=-1)]
-
-    @torch.no_grad
-    def unproject(self, tensors, var, current):
-        return vec_to_tensors(vec=tensors[0], reference=var.params)
-
-
-
-class TensorizeProjection(Projection):
+class TensorizeProjection(ProjectionBase):
     """flattens and concatenates all parameters into a vector and then reshapes it into a tensor"""
     def __init__(self, modules: Chainable, max_side: int, project_update=True, project_params=False, project_grad=False):
         defaults = dict(max_side=max_side)
         super().__init__(modules, defaults=defaults, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
     @torch.no_grad
-    def project(self, tensors, var, current):
-        params = var.params
+    def project(self, tensors, params, grads, loss, states, settings, current):
         max_side = self.settings[params[0]]['max_side']
         num_elems = sum(t.numel() for t in tensors)
 
@@ -60,23 +43,23 @@ class TensorizeProjection(Projection):
         return [vec.view(dims)]
 
     @torch.no_grad
-    def unproject(self, tensors, var, current):
+    def unproject(self, projected_tensors, params, grads, loss, states, settings, current):
         remainder = self.global_state['remainder']
         # warnings.warn(f'{tensors[0].shape = }')
-        vec = tensors[0].view(-1)
+        vec = projected_tensors[0].view(-1)
         if remainder > 0: vec = vec[:-remainder]
-        return vec_to_tensors(vec, var.params)
+        return vec_to_tensors(vec, params)
 
-class BlockPartition(Projection):
+class BlockPartition(ProjectionBase):
     """splits parameters into blocks (for now flatttens them and chunks)"""
     def __init__(self, modules: Chainable, max_size: int, batched: bool = False, project_update=True, project_params=False, project_grad=False):
         defaults = dict(max_size=max_size, batched=batched)
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad, defaults=defaults)
 
     @torch.no_grad
-    def project(self, tensors, var, current):
+    def project(self, tensors, params, grads, loss, states, settings, current):
         partitioned = []
-        for p,t in zip(var.params, tensors):
+        for p,t in zip(params, tensors):
             settings = self.settings[p]
             max_size = settings['max_size']
             n = t.numel()
@@ -101,10 +84,10 @@ class BlockPartition(Projection):
         return partitioned
 
     @torch.no_grad
-    def unproject(self, tensors, var, current):
-        ti = iter(tensors)
+    def unproject(self, projected_tensors, params, grads, loss, states, settings, current):
+        ti = iter(projected_tensors)
         unprojected = []
-        for p in var.params:
+        for p in params:
             settings = self.settings[p]
             n = p.numel()
 
@@ -124,28 +107,3 @@ class BlockPartition(Projection):
 
         return unprojected
 
-
-class TensorNormsProjection(Projection):
-    def __init__(self, modules: Chainable, project_update=True, project_params=False, project_grad=False):
-        super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad)
-
-    @torch.no_grad
-    def project(self, tensors, var, current):
-        orig = self.get_state(var.params, f'{current}_orig')
-        torch._foreach_copy_(orig, tensors)
-
-        norms = torch._foreach_norm(tensors)
-        self.get_state(var.params, f'{current}_orig_norms', cls=TensorList).set_(norms)
-
-        return [torch.stack(norms)]
-
-    @torch.no_grad
-    def unproject(self, tensors, var, current):
-        orig = self.get_state(var.params, f'{current}_orig')
-        orig_norms = torch.stack(self.get_state(var.params, f'{current}_orig_norms'))
-        target_norms = tensors[0]
-
-        orig_norms = torch.where(orig_norms == 0, 1, orig_norms)
-
-        torch._foreach_mul_(orig, (target_norms/orig_norms).detach().cpu().tolist())
-        return orig

torchzero/modules/functional.py

@@ -7,10 +7,19 @@ storage is always indicated in the docstring.
 
 Additional functional variants are present in most module files, e.g. `adam_`, `rmsprop_`, `lion_`, etc.
 """
+from collections.abc import Callable
+from typing import overload
 
-from collections.abc import Callable, Sequence
+import torch
 
-from ..utils import NumberList, TensorList
+from ..utils import (
+    NumberList,
+    TensorList,
+    generic_finfo_eps,
+    generic_max,
+    generic_sum,
+    tofloat,
+)
 
 inf = float('inf')
 
@@ -86,10 +95,10 @@ def root(tensors_:TensorList, p:float, inplace: bool):
         if p == 1: return tensors_.abs_()
         if p == 2: return tensors_.sqrt_()
         return tensors_.pow_(1/p)
-    else:
-        if p == 1: return tensors_.abs()
-        if p == 2: return tensors_.sqrt()
-        return tensors_.pow(1/p)
+
+    if p == 1: return tensors_.abs()
+    if p == 2: return tensors_.sqrt()
+    return tensors_.pow(1/p)
 
 
 def ema_(
@@ -206,4 +215,41 @@ def sqrt_centered_ema_sq_(
         ema_sq_fn=lambda *a, **kw: centered_ema_sq_(*a, **kw, exp_avg_=exp_avg_)
     )
 
+def initial_step_size(tensors: torch.Tensor | TensorList, eps=None) -> float:
+    """initial scaling taken from pytorch L-BFGS to avoid requiring a lot of line search iterations,
+    this version is safer and makes sure largest value isn't smaller than epsilon."""
+    tensors_abs = tensors.abs()
+    tensors_sum = generic_sum(tensors_abs)
+    tensors_max = generic_max(tensors_abs)
+
+    feps = generic_finfo_eps(tensors)
+    if eps is None: eps = feps
+    else: eps = max(eps, feps)
+
+    # scale should not make largest value smaller than epsilon
+    min = eps / tensors_max
+    if min >= 1: return 1.0
+
+    scale = 1 / tensors_sum
+    scale = scale.clip(min=min.item(), max=1)
+    return scale.item()
+
+
+def epsilon_step_size(tensors: torch.Tensor | TensorList, alpha=1e-7) -> float:
+    """makes sure largest value isn't smaller than epsilon."""
+    tensors_abs = tensors.abs()
+    tensors_max = generic_max(tensors_abs)
+    if tensors_max < alpha: return 1.0
+
+    if tensors_max < 1: alpha = alpha / tensors_max
+    return tofloat(alpha)
+
+
+
+def safe_clip(x: torch.Tensor, min=None):
+    """makes sure absolute value of scalar tensor x is not smaller than min"""
+    assert x.numel() == 1, x.shape
+    if min is None: min = torch.finfo(x.dtype).tiny * 2
 
+    if x.abs() < min: return x.new_full(x.size(), min).copysign(x)
+    return x