torchzero 0.1.8__py3-none-any.whl → 0.3.2__py3-none-any.whl

torchzero/modules/experimental/subspace_preconditioners.py

@@ -0,0 +1,128 @@
+from collections import deque
+
+import torch
+# import visualbench as vb
+
+# import torchzero as tz
+
+from ...core import Transform, Chainable, apply
+from ...utils.linalg import inv_sqrt_2x2, matrix_power_eigh, gram_schmidt
+from ...utils import TensorList, vec_to_tensors_
+
+
+def inverse_sqrt(M):
+    if M.shape[-1] == 2: return inv_sqrt_2x2(M, force_pd=True) # general formula for 2x2 matrices
+    return matrix_power_eigh(M, -1/2)
+
+def update_subspace_preconditioner_(
+    grad: torch.Tensor, # store grads and basis as vectors for matmul
+    basis: torch.Tensor, # ndim, k
+    accumulator_: torch.Tensor, # k, k
+    beta: float | None,
+):
+    projected = basis.T @ grad # k
+    outer = torch.outer(projected, projected)
+
+    if beta is None: accumulator_.add_(outer)
+    else: accumulator_.lerp_(outer, 1-beta)
+
+def apply_subspace_preconditioner(
+    tensor: torch.Tensor,
+    basis: torch.Tensor, # ndim, k
+    accumulator: torch.Tensor,
+):
+    preconditioner = inverse_sqrt(accumulator) # k,k
+
+    tensor_projected = basis.T @ tensor # k
+    update_projected = preconditioner @ tensor_projected # k
+    return basis @ update_projected # d
+
+class RandomSubspacePreconditioning(Transform):
+    """full matrix rmsprop in random subspace"""
+    def __init__(self, k: int, beta: float | None = 0.99):
+        defaults = dict(k=k, beta=beta)
+        super().__init__(defaults, uses_grad=False)
+
+    def transform(self, tensors, params, grads, vars):
+        settings = self.settings[params[0]]
+        g = torch.cat([t.view(-1) for t in tensors])
+        k = settings['k']
+        beta = settings['beta']
+
+        if 'basis' not in self.global_state:
+            self.global_state['basis'] = torch.randn(g.numel(), k, device=g.device, dtype=g.dtype)
+            self.global_state['accumulator'] = torch.eye(k, device=g.device, dtype=g.dtype)
+
+        basis = self.global_state['basis']
+        accumulator = self.global_state['accumulator']
+
+        update_subspace_preconditioner_(g, basis, accumulator, beta)
+        try:
+            preconditioned = apply_subspace_preconditioner(g, basis, accumulator)
+        except torch.linalg.LinAlgError:
+            denom = g.abs().sum()
+            if denom <= 1e-10: denom = torch.ones_like(denom)
+            preconditioned = g / g.abs().sum()
+        vec_to_tensors_(preconditioned, tensors)
+
+        return tensors
+
+
+class HistorySubspacePreconditioning(Transform):
+    """full matrix rmsprop in subspace spanned by history of gradient differences
+
+    basis_beta is how much basis is allowed to change, and beta is for preconditioner itself in the basis.
+    """
+    def __init__(self, k: int, beta: float | None = 0.99, basis_beta=0.99, inner: Chainable | None = None):
+        defaults = dict(k=k, beta=beta, basis_beta=basis_beta)
+        super().__init__(defaults, uses_grad=False)
+
+        if inner is not None: self.set_child('inner', inner)
+
+    def transform(self, tensors, params, grads, vars):
+        settings = self.settings[params[0]]
+
+        g = torch.cat([t.view(-1) for t in tensors])
+        k = settings['k']
+        beta = settings['beta']
+        basis_beta = settings['basis_beta']
+
+        if 'history' not in self.global_state:
+            self.global_state['history'] = deque(maxlen=k)
+            self.global_state['accumulator'] = torch.eye(k, device=g.device, dtype=g.dtype)
+            self.global_state['basis'] = torch.ones(g.numel(), k, device=g.device, dtype=g.dtype)
+
+
+        history: deque = self.global_state['history']
+        accumulator = self.global_state['accumulator']
+        basis = self.global_state['basis']
+
+        history.append(g)
+        if len(history) < k:
+            basis_t = torch.randn(g.numel(), k, device=g.device, dtype=g.dtype)
+            history_basis = torch.stack(tuple(history), -1)
+            basis_t[:, -len(history):] = history_basis
+
+        else:
+            basis_t = torch.stack(tuple(history), -1)
+
+        basis_t[:,:-1] = basis_t[:, :-1] - basis_t[:, 1:]
+        basis_t = (basis_t - basis_t.mean()) / basis_t.std()
+
+        basis.lerp_(basis_t, 1-basis_beta)
+        update_subspace_preconditioner_(g, basis, accumulator, beta)
+
+        if 'inner' in self.children:
+            tensors = apply(self.children['inner'], tensors, params, grads, vars)
+            g = torch.cat([t.view(-1) for t in tensors])
+
+        try:
+            preconditioned = apply_subspace_preconditioner(g, basis, accumulator)
+        except torch.linalg.LinAlgError:
+            denom = g.abs().sum()
+            if denom <= 1e-10: denom = torch.ones_like(denom)
+            preconditioned = g / g.abs().sum()
+        vec_to_tensors_(preconditioned, tensors)
+
+        return tensors
+

torchzero/modules/experimental/tropical_newton.py

@@ -0,0 +1,136 @@
+import warnings
+from functools import partial
+from typing import Literal
+from collections.abc import Callable
+import torch
+
+from ...core import Chainable, apply, Module
+from ...utils import vec_to_tensors, TensorList
+from ...utils.derivatives import (
+    hessian_list_to_mat,
+    hessian_mat,
+    jacobian_and_hessian_wrt,
+)
+from ..second_order.newton import lu_solve, cholesky_solve, least_squares_solve
+
+def tropical_sum(x, dim): return torch.amax(x, dim=dim)
+def tropical_mul(x, y): return x+y
+
+def tropical_matmul(x: torch.Tensor, y: torch.Tensor):
+    # this imlements matmul by calling mul and sum
+
+    x_squeeze = False
+    y_squeeze = False
+
+    if x.ndim == 1:
+        x_squeeze = True
+        x = x.unsqueeze(0)
+
+    if y.ndim == 1:
+        y_squeeze = True
+        y = y.unsqueeze(1)
+
+    res = tropical_sum(tropical_mul(x.unsqueeze(-1), y.unsqueeze(-3)), dim = -2)
+
+    if x_squeeze: res = res.squeeze(-2)
+    if y_squeeze: res = res.squeeze(-1)
+
+    return res
+
+def tropical_dot(x:torch.Tensor, y:torch.Tensor):
+    assert x.ndim == 1 and y.ndim == 1
+    return tropical_matmul(x.unsqueeze(0), y.unsqueeze(1))
+
+def tropical_outer(x:torch.Tensor, y:torch.Tensor):
+    assert x.ndim == 1 and y.ndim == 1
+    return tropical_matmul(x.unsqueeze(1), y.unsqueeze(0))
+
+
+def tropical_solve(A: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    r = b.unsqueeze(1) - A
+    return r.amin(dim=-2)
+
+def tropical_solve_and_reconstruct(A: torch.Tensor, b: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    r = b.unsqueeze(1) - A
+    x = r.amin(dim=-2)
+    b_hat = tropical_matmul(A, x)
+    return x, b_hat
+
+def tikhonov(H: torch.Tensor, reg: float):
+    if reg!=0: H += torch.eye(H.size(-1), dtype=H.dtype, device=H.device) * reg
+    return H
+
+
+class TropicalNewton(Module):
+    """suston"""
+    def __init__(
+        self,
+        reg: float | None = None,
+        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
+        vectorize: bool = True,
+        interpolate:bool=False,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(reg=reg, hessian_method=hessian_method, vectorize=vectorize, interpolate=interpolate)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+        closure = vars.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        reg = settings['reg']
+        hessian_method = settings['hessian_method']
+        vectorize = settings['vectorize']
+        interpolate = settings['interpolate']
+
+        # ------------------------ calculate grad and hessian ------------------------ #
+        if hessian_method == 'autograd':
+            with torch.enable_grad():
+                loss = vars.loss = vars.loss_approx = closure(False)
+                g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
+                g_list = [t[0] for t in g_list] # remove leading dim from loss
+                vars.grad = g_list
+                H = hessian_list_to_mat(H_list)
+
+        elif hessian_method in ('func', 'autograd.functional'):
+            strat = 'forward-mode' if vectorize else 'reverse-mode'
+            with torch.enable_grad():
+                g_list = vars.get_grad(retain_graph=True)
+                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)
+
+        # -------------------------------- inner step -------------------------------- #
+        if 'inner' in self.children:
+            g_list = apply(self.children['inner'], list(g_list), params=params, grads=list(g_list), vars=vars)
+        g = torch.cat([t.view(-1) for t in g_list])
+
+        # ------------------------------- regulazition ------------------------------- #
+        if reg is not None: H = tikhonov(H, reg)
+
+        # ----------------------------------- solve ---------------------------------- #
+        tropical_update, g_hat = tropical_solve_and_reconstruct(H, g)
+
+        g_norm = torch.linalg.vector_norm(g) # pylint:disable=not-callable
+        abs_error = torch.linalg.vector_norm(g-g_hat) # pylint:disable=not-callable
+        rel_error = abs_error/g_norm.clip(min=1e-8)
+
+        if interpolate:
+            if rel_error > 1e-8:
+
+                update = cholesky_solve(H, g)
+                if update is None: update = lu_solve(H, g)
+                if update is None: update = least_squares_solve(H, g)
+
+                tropical_update.lerp_(update.ravel(), rel_error.clip(max=1))
+
+        vars.update = vec_to_tensors(tropical_update, params)
+        return vars

torchzero/modules/functional.py

@@ -0,0 +1,209 @@
+"""
+Arguments that are modified in-place are denoted with "_" at the end.
+
+Some functions return one of the arguments which was modified in-place, some return new tensors.
+Make sure to keep track of that to avoid unexpected in-place modifications of buffers. The returned
+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, Sequence
+
+from ..utils import NumberList, TensorList
+
+inf = float('inf')
+
+def debiased_step_size(
+    step,
+    beta1: float | NumberList | None = None,
+    beta2: float | NumberList | None = None,
+    pow: float = 2,
+    alpha: float | NumberList = 1,
+):
+    """returns multiplier to step size"""
+    if isinstance(beta1, NumberList): beta1 = beta1.fill_none(0)
+    if isinstance(beta2, NumberList): beta2 = beta2.fill_none(0)
+
+    step_size = alpha
+    if beta1 is not None:
+        bias_correction1 = 1.0 - (beta1 ** step)
+        step_size /= bias_correction1
+    if beta2 is not None:
+        bias_correction2 = 1.0 - (beta2 ** step)
+        step_size *= bias_correction2 ** (1/pow)
+    return step_size
+
+def debias(
+    tensors_: TensorList,
+    step: int,
+    inplace: bool,
+    beta1: float | NumberList | None = None,
+    beta2: float | NumberList | None = None,
+    alpha: float | NumberList = 1,
+    pow: float = 2,
+):
+    step_size = debiased_step_size(step=step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
+    if inplace: return tensors_.mul_(step_size)
+    return tensors_ * step_size
+
+def debias_second_momentum(tensors_:TensorList, step: int, beta: float | NumberList, pow: float, inplace:bool):
+    """debias 2nd momentum, optionally in-place"""
+    bias_correction2 = (1.0 - (beta ** step)) ** (1/pow)
+    if inplace: return tensors_.div_(bias_correction2)
+    return tensors_ / bias_correction2
+
+def lerp_power_(tensors:TensorList, exp_avg_pow_:TensorList, beta:float|NumberList, pow:float) -> TensorList:
+    """
+    Lerp `exp_avg_pow_` with `tensors ^ pow`
+
+    Returns `exp_avg_pow_`.
+    """
+    if pow == 1: return exp_avg_pow_.lerp_(tensors.abs(), 1-beta)
+    if pow == 2: return exp_avg_pow_.mul_(beta).addcmul_(tensors, tensors, value = 1-beta)
+    if pow % 2 == 0: return exp_avg_pow_.lerp_(tensors.pow(pow), 1-beta)
+    return exp_avg_pow_.lerp_(tensors.pow(pow).abs_(), 1-beta)
+
+def add_power_(tensors:TensorList, sum_:TensorList, pow:float) -> TensorList:
+    """
+    Add `tensors ^ pow` to `sum_`
+
+    Returns `sum_`.
+    """
+    if pow == 1: return sum_.add_(tensors.abs())
+    if pow == 2: return sum_.addcmul_(tensors, tensors)
+    if pow % 2 == 0: return sum_.add_(tensors.pow(pow))
+    return sum_.add_(tensors.pow(pow).abs_())
+
+
+def root(tensors_:TensorList, p:float, inplace: bool):
+    """
+    Root of tensors, optionally in-place.
+
+    Returns `tensors_` if `inplace` else new tensors.
+    """
+    if inplace:
+        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)
+
+
+def ema_(
+    tensors: TensorList,
+    exp_avg_: TensorList,
+    beta: float | NumberList,
+    dampening: float | NumberList = 0,
+    lerp: bool = True,
+):
+    """
+    Updates `exp_avg_` with EMA of `tensors`.
+
+    Returns `exp_avg_`.
+    """
+    tensors.lazy_mul_(1 - dampening)
+    if lerp: return exp_avg_.lerp_(tensors, (1 - beta))
+    return exp_avg_.mul_(beta).add_(tensors)
+
+def ema_sq_(
+    tensors: TensorList,
+    exp_avg_sq_: TensorList,
+    beta: float | NumberList,
+    max_exp_avg_sq_: TensorList | None,
+    pow: float = 2,
+):
+    """
+    Updates `exp_avg_sq_` with EMA of squared `tensors`, if `max_exp_avg_sq_` is not None, updates it with maximum of EMA.
+
+    Returns `exp_avg_sq_` or `max_exp_avg_sq_`.
+    """
+    lerp_power_(tensors=tensors, exp_avg_pow_=exp_avg_sq_,beta=beta,pow=pow)
+
+    # AMSGrad
+    if max_exp_avg_sq_ is not None:
+        max_exp_avg_sq_.maximum_(exp_avg_sq_)
+        exp_avg_sq_ = max_exp_avg_sq_
+
+    return exp_avg_sq_
+
+def sqrt_ema_sq_(
+    tensors: TensorList,
+    exp_avg_sq_: TensorList,
+    beta: float | NumberList,
+    max_exp_avg_sq_: TensorList | None,
+    debiased: bool,
+    step: int,
+    pow: float = 2,
+    ema_sq_fn: Callable = ema_sq_,
+):
+    """
+    Updates `exp_avg_sq_` with EMA of squared `tensors` and calculates it's square root,
+    with optional AMSGrad and debiasing.
+
+    Returns new tensors.
+    """
+    exp_avg_sq_=ema_sq_fn(
+        tensors=tensors,
+        exp_avg_sq_=exp_avg_sq_,
+        beta=beta,
+        max_exp_avg_sq_=max_exp_avg_sq_,
+        pow=pow,
+    )
+
+    sqrt_exp_avg_sq = root(exp_avg_sq_, pow, inplace=False)
+
+    if debiased: sqrt_exp_avg_sq = debias_second_momentum(sqrt_exp_avg_sq, step=step, beta=beta, pow=pow, inplace=True)
+    return sqrt_exp_avg_sq
+
+
+def centered_ema_sq_(tensors: TensorList, exp_avg_: TensorList, exp_avg_sq_: TensorList,
+                     beta: float | NumberList, max_exp_avg_sq_: TensorList | None = None, pow:float=2):
+    """
+    Updates `exp_avg_` and `exp_avg_sq_` with EMA of `tensors` and squared `tensors`,
+    centers `exp_avg_sq_` by subtracting `exp_avg_` squared.
+
+    Returns `max_exp_avg_sq_` or new tensors.
+    """
+    exp_avg_sq_ = ema_sq_(tensors, exp_avg_sq_=exp_avg_sq_, beta=beta, max_exp_avg_sq_=max_exp_avg_sq_, pow=pow)
+    exp_avg_.lerp_(tensors, 1-beta)
+    exp_avg_sq_ = exp_avg_sq_.addcmul(exp_avg_, exp_avg_, value=-1)
+
+    # AMSGrad
+    if max_exp_avg_sq_ is not None:
+        max_exp_avg_sq_.maximum_(exp_avg_sq_)
+        exp_avg_sq_ = max_exp_avg_sq_
+
+    return exp_avg_sq_
+
+def sqrt_centered_ema_sq_(
+    tensors: TensorList,
+    exp_avg_: TensorList,
+    exp_avg_sq_: TensorList,
+    max_exp_avg_sq_: TensorList | None,
+    beta: float | NumberList,
+    debiased: bool,
+    step: int,
+    pow: float = 2,
+):
+    """
+    Updates `exp_avg_` and `exp_avg_sq_` with EMA of `tensors` and squared `tensors`,
+    centers `exp_avg_sq_` by subtracting `exp_avg_` squared. Calculates it's square root,
+    with optional AMSGrad and debiasing.
+
+    Returns new tensors.
+    """
+    return sqrt_ema_sq_(
+        tensors=tensors,
+        exp_avg_sq_=exp_avg_sq_,
+        beta=beta,
+        max_exp_avg_sq_=max_exp_avg_sq_,
+        debiased=debiased,
+        step=step,
+        pow=pow,
+        ema_sq_fn=lambda *a, **kw: centered_ema_sq_(*a, **kw, exp_avg_=exp_avg_)
+    )
+
+

torchzero/modules/grad_approximation/__init__.py

@@ -0,0 +1,4 @@
+from .grad_approximator import GradApproximator, GradTarget
+from .fdm import FDM
+from .rfdm import RandomizedFDM, MeZO, SPSA, RDSA, GaussianSmoothing
+from .forward_gradient import ForwardGradient

torchzero/modules/grad_approximation/fdm.py

@@ -0,0 +1,120 @@
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+from ...utils import TensorList
+from .grad_approximator import GradApproximator, GradTarget, _FD_Formula
+
+
+def _forward2(closure: Callable[..., float], param:torch.Tensor, idx: int, h, v_0: float | None):
+    if v_0 is None: v_0 = closure(False)
+    assert param.ndim == 1
+    param[idx] += h
+    v_plus = closure(False)
+    param[idx] -= h
+    return v_0, v_0, (v_plus - v_0) / h # (loss, loss_approx, grad)
+
+def _backward2(closure: Callable[..., float], param:torch.Tensor, idx: int, h, v_0: float | None):
+    if v_0 is None: v_0 = closure(False)
+    assert param.ndim == 1
+    param[idx] -= h
+    v_minus = closure(False)
+    param[idx] += h
+    return v_0, v_0, (v_0 - v_minus) / h
+
+def _central2(closure: Callable[..., float], param:torch.Tensor, idx: int, h, v_0: Any):
+    assert param.ndim == 1
+    param[idx] += h
+    v_plus = closure(False)
+
+    param[idx] -= h * 2
+    v_minus = closure(False)
+
+    param[idx] += h
+    return v_0, v_plus, (v_plus - v_minus) / (2 * h)
+
+def _forward3(closure: Callable[..., float], param:torch.Tensor, idx: int, h, v_0: float | None):
+    if v_0 is None: v_0 = closure(False)
+    assert param.ndim == 1
+    param[idx] += h
+    v_plus1 = closure(False)
+
+    param[idx] += h
+    v_plus2 = closure(False)
+
+    param[idx] -= 2 * h
+    return v_0, v_0, (-3*v_0 + 4*v_plus1 - v_plus2) / (2 * h)
+
+def _backward3(closure: Callable[..., float], param:torch.Tensor, idx: int, h, v_0: float | None):
+    if v_0 is None: v_0 = closure(False)
+    assert param.ndim == 1
+    param[idx] -= h
+    v_minus1 = closure(False)
+
+    param[idx] -= h
+    v_minus2 = closure(False)
+
+    param[idx] += 2 * h
+    return v_0, v_0, (v_minus2 - 4*v_minus1 + 3*v_0) / (2 * h)
+
+def _central4(closure: Callable[..., float], param:torch.Tensor, idx: int, h, v_0: Any):
+    assert param.ndim == 1
+
+    param[idx] += h
+    v_plus1 = closure(False)
+
+    param[idx] += h
+    v_plus2 = closure(False)
+
+    param[idx] -= 3 * h
+    v_minus1 = closure(False)
+
+    param[idx] -= h
+    v_minus2 = closure(False)
+
+    param[idx] += 2 * h
+    return v_0, v_plus1, (v_minus2 - 8*v_minus1 + 8*v_plus1 - v_plus2) / (12 * h)
+
+_FD_FUNCS = {
+    "forward2": _forward2,
+    "backward2": _backward2,
+    "central2": _central2,
+    "central3": _central2, # they are the same
+    "forward3": _forward3,
+    "backward3": _backward3,
+    "central4": _central4,
+}
+
+
+class FDM(GradApproximator):
+    """Approximate gradients via finite difference method
+
+    Args:
+        h (float, optional): magnitude of parameter perturbation. Defaults to 1e-3.
+        formula (_FD_Formula, optional): finite difference formula. Defaults to 'central2'.
+        target (GradTarget, optional): what to set on vars. Defaults to 'closure'.
+    """
+    def __init__(self, h: float=1e-3, formula: _FD_Formula = 'central2', target: GradTarget = 'closure'):
+        defaults = dict(h=h, formula=formula)
+        super().__init__(defaults, target=target)
+
+    @torch.no_grad
+    def approximate(self, closure, params, loss, vars):
+        grads = []
+        loss_approx = None
+
+        for p in params:
+            g = torch.zeros_like(p)
+            grads.append(g)
+
+            settings = self.settings[p]
+            h = settings['h']
+            fd_fn = _FD_FUNCS[settings['formula']]
+
+            p_flat = p.view(-1); g_flat = g.view(-1)
+            for i in range(len(p_flat)):
+                loss, loss_approx, d = fd_fn(closure=closure, param=p_flat, idx=i, h=h, v_0=loss)
+                g_flat[i] = d
+
+        return grads, loss, loss_approx

torchzero/modules/grad_approximation/forward_gradient.py

@@ -0,0 +1,81 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import torch
+
+from ...utils import Distributions, NumberList, TensorList, generic_eq
+from ...utils.derivatives import jvp, jvp_fd_central, jvp_fd_forward
+from .grad_approximator import GradApproximator, GradTarget
+from .rfdm import RandomizedFDM
+
+
+class ForwardGradient(RandomizedFDM):
+    """Forward gradient method, same as randomized finite difference but directional derivative is estimated via autograd (as jacobian vector product)
+
+    Args:
+        n_samples (int, optional): number of random gradient samples. Defaults to 1.
+        distribution (Distributions, optional): distribution for random gradient samples. Defaults to "gaussian".
+        beta (float, optional):
+            if not 0, acts as momentum on gradient samples, making the subspace spanned by them change slowly. Defaults to 0.
+        pre_generate (bool, optional):
+            whether to pre-generate gradient samples before each step. Defaults to True.
+        jvp_method (str, optional):
+            how to calculate jacobian vector product, note that with `forward` and 'central' this is identical to randomized finite difference. Defaults to 'autograd'.
+        h (float, optional): finite difference step size of jvp_method is set to `forward` or `central`. Defaults to 1e-3.
+        target (GradTarget, optional): what to set on vars. Defaults to "closure".
+    """
+    PRE_MULTIPLY_BY_H = False
+    def __init__(
+        self,
+        n_samples: int = 1,
+        distribution: Distributions = "gaussian",
+        beta: float = 0,
+        pre_generate = True,
+        jvp_method: Literal['autograd', 'forward', 'central'] = 'autograd',
+        h: float = 1e-3,
+        target: GradTarget = "closure",
+        seed: int | None | torch.Generator = None,
+    ):
+        super().__init__(h=h, n_samples=n_samples, distribution=distribution, beta=beta, target=target, pre_generate=pre_generate, seed=seed)
+        self.defaults['jvp_method'] = jvp_method
+
+    @torch.no_grad
+    def approximate(self, closure, params, loss, vars):
+        params = TensorList(params)
+        loss_approx = None
+
+        settings = self.settings[params[0]]
+        n_samples = settings['n_samples']
+        jvp_method = settings['jvp_method']
+        h = settings['h']
+        distribution = settings['distribution']
+        default = [None]*n_samples
+        perturbations = list(zip(*(self.state[p].get('perturbations', default) for p in params)))
+        generator = self._get_generator(settings['seed'], params)
+
+        grad = None
+        for i in range(n_samples):
+            prt = perturbations[i]
+            if prt[0] is None: prt = params.sample_like(distribution=distribution, generator=generator)
+            else: prt = TensorList(prt)
+
+            if jvp_method == 'autograd':
+                with torch.enable_grad():
+                    loss, d = jvp(partial(closure, False), params=params, tangent=prt)
+
+            elif jvp_method == 'forward':
+                loss, d = jvp_fd_forward(partial(closure, False), params=params, tangent=prt, v_0=loss, normalize=True, h=h)
+
+            elif jvp_method == 'central':
+                loss_approx, d = jvp_fd_central(partial(closure, False), params=params, tangent=prt, normalize=True, h=h)
+
+            else: raise ValueError(jvp_method)
+
+            if grad is None: grad = prt * d
+            else: grad += prt * d
+
+        assert grad is not None
+        if n_samples > 1: grad.div_(n_samples)
+        return grad, loss, loss_approx
+