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

torchzero/modules/second_order/nystrom.py

@@ -6,7 +6,7 @@ import torch
 from ...utils import TensorList, as_tensorlist, generic_zeros_like, generic_vector_norm, generic_numel, vec_to_tensors
 from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
 
-from ...core import Chainable, apply, Module
+from ...core import Chainable, apply_transform, Module
 from ...utils.linalg.solve import nystrom_sketch_and_solve, nystrom_pcg
 
 class NystromSketchAndSolve(Module):
@@ -15,7 +15,7 @@ class NystromSketchAndSolve(Module):
         rank: int,
         reg: float = 1e-3,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
-        h=1e-2,
+        h=1e-3,
         inner: Chainable | None = None,
         seed: int | None = None,
     ):
@@ -26,10 +26,10 @@ class NystromSketchAndSolve(Module):
             self.set_child('inner', inner)
 
     @torch.no_grad
-    def step(self, vars):
-        params = TensorList(vars.params)
+    def step(self, var):
+        params = TensorList(var.params)
 
-        closure = vars.closure
+        closure = var.closure
         if closure is None: raise RuntimeError('NewtonCG requires closure')
 
         settings = self.settings[params[0]]
@@ -47,7 +47,7 @@ class NystromSketchAndSolve(Module):
 
         # ---------------------- Hessian vector product function --------------------- #
         if hvp_method == 'autograd':
-            grad = vars.get_grad(create_graph=True)
+            grad = var.get_grad(create_graph=True)
 
             def H_mm(x):
                 with torch.enable_grad():
@@ -57,7 +57,7 @@ class NystromSketchAndSolve(Module):
         else:
 
             with torch.enable_grad():
-                grad = vars.get_grad()
+                grad = var.get_grad()
 
             if hvp_method == 'forward':
                 def H_mm(x):
@@ -74,14 +74,14 @@ class NystromSketchAndSolve(Module):
 
 
         # -------------------------------- inner step -------------------------------- #
-        b = vars.get_update()
+        b = var.get_update()
         if 'inner' in self.children:
-            b = apply(self.children['inner'], b, params=params, grads=grad, vars=vars)
+            b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
 
         # ------------------------------ sketch&n&solve ------------------------------ #
         x = nystrom_sketch_and_solve(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), rank=rank, reg=reg, generator=generator)
-        vars.update = vec_to_tensors(x, reference=params)
-        return vars
+        var.update = vec_to_tensors(x, reference=params)
+        return var
 
 
 
@@ -93,7 +93,7 @@ class NystromPCG(Module):
         tol=1e-3,
         reg: float = 1e-6,
         hvp_method: Literal["forward", "central", "autograd"] = "autograd",
-        h=1e-2,
+        h=1e-3,
         inner: Chainable | None = None,
         seed: int | None = None,
     ):
@@ -104,10 +104,10 @@ class NystromPCG(Module):
             self.set_child('inner', inner)
 
     @torch.no_grad
-    def step(self, vars):
-        params = TensorList(vars.params)
+    def step(self, var):
+        params = TensorList(var.params)
 
-        closure = vars.closure
+        closure = var.closure
         if closure is None: raise RuntimeError('NewtonCG requires closure')
 
         settings = self.settings[params[0]]
@@ -129,7 +129,7 @@ class NystromPCG(Module):
 
         # ---------------------- Hessian vector product function --------------------- #
         if hvp_method == 'autograd':
-            grad = vars.get_grad(create_graph=True)
+            grad = var.get_grad(create_graph=True)
 
             def H_mm(x):
                 with torch.enable_grad():
@@ -139,7 +139,7 @@ class NystromPCG(Module):
         else:
 
             with torch.enable_grad():
-                grad = vars.get_grad()
+                grad = var.get_grad()
 
             if hvp_method == 'forward':
                 def H_mm(x):
@@ -156,13 +156,13 @@ class NystromPCG(Module):
 
 
         # -------------------------------- inner step -------------------------------- #
-        b = vars.get_update()
+        b = var.get_update()
         if 'inner' in self.children:
-            b = apply(self.children['inner'], b, params=params, grads=grad, vars=vars)
+            b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
 
         # ------------------------------ sketch&n&solve ------------------------------ #
         x = nystrom_pcg(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), sketch_size=sketch_size, reg=reg, tol=tol, maxiter=maxiter, x0_=None, generator=generator)
-        vars.update = vec_to_tensors(x, reference=params)
-        return vars
+        var.update = vec_to_tensors(x, reference=params)
+        return var
 
 

torchzero/modules/smoothing/gaussian.py

@@ -6,7 +6,7 @@ from typing import Literal
 
 import torch
 
-from ...core import Modular, Module, Vars
+from ...core import Modular, Module, Var
 from ...utils import NumberList, TensorList
 from ...utils.derivatives import jacobian_wrt
 from ..grad_approximation import GradApproximator, GradTarget
@@ -17,24 +17,24 @@ class Reformulation(Module, ABC):
         super().__init__(defaults)
 
     @abstractmethod
-    def closure(self, backward: bool, closure: Callable, params:list[torch.Tensor], vars: Vars) -> tuple[float | torch.Tensor, Sequence[torch.Tensor] | None]:
+    def closure(self, backward: bool, closure: Callable, params:list[torch.Tensor], var: Var) -> tuple[float | torch.Tensor, Sequence[torch.Tensor] | None]:
         """returns loss and gradient, if backward is False then gradient can be None"""
 
-    def pre_step(self, vars: Vars) -> Vars | None:
+    def pre_step(self, var: Var) -> Var | None:
         """This runs once before each step, whereas `closure` may run multiple times per step if further modules
         evaluate gradients at multiple points. This is useful for example to pre-generate new random perturbations."""
-        return vars
+        return var
 
-    def step(self, vars):
-        ret = self.pre_step(vars)
-        if isinstance(ret, Vars): vars = ret
+    def step(self, var):
+        ret = self.pre_step(var)
+        if isinstance(ret, Var): var = ret
 
-        if vars.closure is None: raise RuntimeError("Reformulation requires closure")
-        params, closure = vars.params, vars.closure
+        if var.closure is None: raise RuntimeError("Reformulation requires closure")
+        params, closure = var.params, var.closure
 
 
         def modified_closure(backward=True):
-            loss, grad = self.closure(backward, closure, params, vars)
+            loss, grad = self.closure(backward, closure, params, var)
 
             if grad is not None:
                 for p,g in zip(params, grad):
@@ -42,8 +42,8 @@ class Reformulation(Module, ABC):
 
             return loss
 
-        vars.closure = modified_closure
-        return vars
+        var.closure = modified_closure
+        return var
 
 
 def _decay_sigma_(self: Module, params):
@@ -58,7 +58,7 @@ def _generate_perturbations_to_state_(self: Module, params: TensorList, n_sample
     for param, prt in zip(params, zip(*perturbations)):
         self.state[param]['perturbations'] = prt
 
-def _clear_state_hook(optimizer: Modular, vars: Vars, self: Module):
+def _clear_state_hook(optimizer: Modular, var: Var, self: Module):
     for m in optimizer.unrolled_modules:
         if m is not self:
             m.reset()
@@ -85,12 +85,12 @@ class GaussianHomotopy(Reformulation):
             else: self.global_state['generator'] = None
         return self.global_state['generator']
 
-    def pre_step(self, vars):
-        params = TensorList(vars.params)
+    def pre_step(self, var):
+        params = TensorList(var.params)
         settings = self.settings[params[0]]
         n_samples = settings['n_samples']
-        init_sigma = self.get_settings('init_sigma', params=params)
-        sigmas = self.get_state('sigma', params = params, init=init_sigma)
+        init_sigma = [self.settings[p]['init_sigma'] for p in params]
+        sigmas = self.get_state(params, 'sigma', init=init_sigma)
 
         if any('perturbations' not in self.state[p] for p in params):
             generator = self._get_generator(settings['seed'], params)
@@ -109,9 +109,9 @@ class GaussianHomotopy(Reformulation):
         tol = settings['tol']
         if tol is not None and not decayed:
             if not any('prev_params' in self.state[p] for p in params):
-                prev_params = self.get_state('prev_params', params=params, cls=TensorList, init='param')
+                prev_params = self.get_state(params, 'prev_params', cls=TensorList, init='param')
             else:
-                prev_params = self.get_state('prev_params', params=params, cls=TensorList, init='param')
+                prev_params = self.get_state(params, 'prev_params', cls=TensorList, init='param')
                 s = params - prev_params
 
                 if s.abs().global_max() <= tol:
@@ -124,10 +124,10 @@ class GaussianHomotopy(Reformulation):
             generator = self._get_generator(settings['seed'], params)
             _generate_perturbations_to_state_(self, params=params, n_samples=n_samples, sigmas=sigmas, generator=generator)
             if settings['clear_state']:
-                vars.post_step_hooks.append(partial(_clear_state_hook, self=self))
+                var.post_step_hooks.append(partial(_clear_state_hook, self=self))
 
     @torch.no_grad
-    def closure(self, backward, closure, params, vars):
+    def closure(self, backward, closure, params, var):
         params = TensorList(params)
 
         settings = self.settings[params[0]]

torchzero/modules/smoothing/laplacian.py

@@ -67,7 +67,7 @@ class LaplacianSmoothing(Transform):
             minimum number of elements in a parameter to apply laplacian smoothing to.
             Only has effect if `layerwise` is True. Defaults to 4.
         target (str, optional):
-            what to set on vars.
+            what to set on var.
 
     Reference:
         *Osher, S., Wang, B., Yin, P., Luo, X., Barekat, F., Pham, M., & Lin, A. (2022).
@@ -82,19 +82,17 @@ class LaplacianSmoothing(Transform):
 
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        layerwise = self.settings[params[0]]['layerwise']
+    def apply(self, tensors, params, grads, loss, states, settings):
+        layerwise = settings[0]['layerwise']
 
         # layerwise laplacian smoothing
         if layerwise:
 
             # precompute the denominator for each layer and store it in each parameters state
             smoothed_target = TensorList()
-            for p, t in zip(params, tensors):
-                settings = self.settings[p]
-                if p.numel() > settings['min_numel']:
-                    state = self.state[p]
-                    if 'denominator' not in state: state['denominator'] = _precompute_denominator(p, settings['sigma'])
+            for p, t, state, setting in zip(params, tensors, states, settings):
+                if p.numel() > setting['min_numel']:
+                    if 'denominator' not in state: state['denominator'] = _precompute_denominator(p, setting['sigma'])
                     smoothed_target.append(torch.fft.ifft(torch.fft.fft(t.view(-1)) / state['denominator']).real.view_as(t)) #pylint:disable=not-callable
                 else:
                     smoothed_target.append(t)
@@ -106,7 +104,7 @@ class LaplacianSmoothing(Transform):
         # precompute full denominator
         tensors = TensorList(tensors)
         if self.global_state.get('full_denominator', None) is None:
-            self.global_state['full_denominator'] = _precompute_denominator(tensors.to_vec(), self.settings[params[0]]['sigma'])
+            self.global_state['full_denominator'] = _precompute_denominator(tensors.to_vec(), settings[0]['sigma'])
 
         # apply the smoothing
         vec = tensors.to_vec()

torchzero/modules/weight_decay/__init__.py

@@ -1 +1 @@
-from .weight_decay import WeightDecay, DirectWeightDecay, decay_weights_
+from .weight_decay import WeightDecay, DirectWeightDecay, decay_weights_, NormalizedWeightDecay

torchzero/modules/weight_decay/weight_decay.py

@@ -1,9 +1,11 @@
 from collections.abc import Iterable, Sequence
+from typing import Literal
 
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList, as_tensorlist
+from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states
+
 
 @torch.no_grad
 def weight_decay_(
@@ -25,12 +27,44 @@ class WeightDecay(Transform):
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        weight_decay = self.get_settings('weight_decay', params=params, cls=NumberList)
-        ord = self.settings[params[0]]['ord']
+    def apply(self, tensors, params, grads, loss, states, settings):
+        weight_decay = NumberList(s['weight_decay'] for s in settings)
+        ord = settings[0]['ord']
 
         return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay, ord)
 
+class NormalizedWeightDecay(Transform):
+    def __init__(
+        self,
+        weight_decay: float = 0.1,
+        ord: int = 2,
+        norm_input: Literal["update", "grad", "params"] = "update",
+        target: Target = "update",
+    ):
+        defaults = dict(weight_decay=weight_decay, ord=ord, norm_input=norm_input)
+        super().__init__(defaults, uses_grad=norm_input == 'grad', target=target)
+
+    @torch.no_grad
+    def apply(self, tensors, params, grads, loss, states, settings):
+        weight_decay = NumberList(s['weight_decay'] for s in settings)
+
+        ord = settings[0]['ord']
+        norm_input = settings[0]['norm_input']
+
+        if norm_input == 'update': src = TensorList(tensors)
+        elif norm_input == 'grad':
+            assert grads is not None
+            src = TensorList(grads)
+        elif norm_input == 'params':
+            src = TensorList(params)
+        else:
+            raise ValueError(norm_input)
+
+        norm = src.global_vector_norm(ord)
+
+        return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay * norm, ord)
+
+
 @torch.no_grad
 def decay_weights_(params: Iterable[torch.Tensor], weight_decay: float | NumberList, ord:int=2):
     """directly decays weights in-place"""
@@ -44,9 +78,9 @@ class DirectWeightDecay(Module):
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, vars):
-        weight_decay = self.get_settings('weight_decay', params=vars.params, cls=NumberList)
-        ord = self.settings[vars.params[0]]['ord']
+    def step(self, var):
+        weight_decay = self.get_settings(var.params, 'weight_decay', cls=NumberList)
+        ord = self.settings[var.params[0]]['ord']
 
-        decay_weights_(vars.params, weight_decay, ord)
-        return vars
+        decay_weights_(var.params, weight_decay, ord)
+        return var

torchzero/modules/wrappers/optim_wrapper.py

@@ -24,8 +24,8 @@ class Wrap(Module):
         return super().set_param_groups(param_groups)
 
     @torch.no_grad
-    def step(self, vars):
-        params = vars.params
+    def step(self, var):
+        params = var.params
 
         # initialize opt on 1st step
         if self.optimizer is None:
@@ -35,18 +35,18 @@ class Wrap(Module):
 
         # set grad to update
         orig_grad = [p.grad for p in params]
-        for p, u in zip(params, vars.get_update()):
+        for p, u in zip(params, var.get_update()):
             p.grad = u
 
         # if this module is last, can step with _opt directly
         # direct step can't be applied if next module is LR but _opt doesn't support lr,
         # and if there are multiple different per-parameter lrs (would be annoying to support)
-        if vars.is_last and (
-            (vars.last_module_lrs is None)
+        if var.is_last and (
+            (var.last_module_lrs is None)
             or
-            (('lr' in self.optimizer.defaults) and (len(set(vars.last_module_lrs)) == 1))
+            (('lr' in self.optimizer.defaults) and (len(set(var.last_module_lrs)) == 1))
         ):
-            lr = 1 if vars.last_module_lrs is None else vars.last_module_lrs[0]
+            lr = 1 if var.last_module_lrs is None else var.last_module_lrs[0]
 
             # update optimizer lr with desired lr
             if lr != 1:
@@ -68,19 +68,19 @@ class Wrap(Module):
             for p, g in zip(params, orig_grad):
                 p.grad = g
 
-            vars.stop = True; vars.skip_update = True
-            return vars
+            var.stop = True; var.skip_update = True
+            return var
 
         # this is not the last module, meaning update is difference in parameters
         params_before_step = [p.clone() for p in params]
         self.optimizer.step() # step and update params
         for p, g in zip(params, orig_grad):
             p.grad = g
-        vars.update = list(torch._foreach_sub(params_before_step, params)) # set update to difference between params
+        var.update = list(torch._foreach_sub(params_before_step, params)) # set update to difference between params
         for p, o in zip(params, params_before_step):
             p.set_(o) # pyright: ignore[reportArgumentType]
 
-        return vars
+        return var
 
     def reset(self):
         super().reset()

torchzero/optim/wrappers/directsearch.py

@@ -0,0 +1,244 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import directsearch
+import numpy as np
+import torch
+from directsearch.ds import DEFAULT_PARAMS
+
+from ...modules.second_order.newton import tikhonov_
+from ...utils import Optimizer, TensorList
+
+
+def _ensure_float(x):
+    if isinstance(x, torch.Tensor): return x.detach().cpu().item()
+    if isinstance(x, np.ndarray): return x.item()
+    return float(x)
+
+def _ensure_numpy(x):
+    if isinstance(x, torch.Tensor): return x.detach().cpu()
+    if isinstance(x, np.ndarray): return x
+    return np.array(x)
+
+
+Closure = Callable[[bool], Any]
+
+
+class DirectSearch(Optimizer):
+    """Use directsearch as pytorch optimizer.
+
+    Note that this performs full minimization on each step,
+    so usually you would want to perform a single step, although performing multiple steps will refine the
+    solution.
+
+    Args:
+        params (_type_): _description_
+        maxevals (_type_, optional): _description_. Defaults to DEFAULT_PARAMS['maxevals'].
+    """
+    def __init__(
+        self,
+        params,
+        maxevals = DEFAULT_PARAMS['maxevals'], # Maximum number of function evaluations
+        rho = DEFAULT_PARAMS['rho'], # Forcing function
+        sketch_dim = DEFAULT_PARAMS['sketch_dim'], # Target dimension for sketching
+        sketch_type = DEFAULT_PARAMS['sketch_type'], # Sketching technique
+        poll_type = DEFAULT_PARAMS['poll_type'], # Polling direction type
+        alpha0 = DEFAULT_PARAMS['alpha0'], # Original stepsize value
+        alpha_max = DEFAULT_PARAMS['alpha_max'], # Maximum value for the stepsize
+        alpha_min = DEFAULT_PARAMS['alpha_min'], # Minimum value for the stepsize
+        gamma_inc = DEFAULT_PARAMS['gamma_inc'], # Increasing factor for the stepsize
+        gamma_dec = DEFAULT_PARAMS['gamma_dec'], # Decreasing factor for the stepsize
+        verbose = DEFAULT_PARAMS['verbose'], # Display information about the method
+        print_freq = DEFAULT_PARAMS['print_freq'], # How frequently to display information
+        use_stochastic_three_points = DEFAULT_PARAMS['use_stochastic_three_points'], # Boolean for a specific method
+        rho_uses_normd = DEFAULT_PARAMS['rho_uses_normd'], # Forcing function based on direction norm
+    ):
+        super().__init__(params, {})
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    def _objective(self, x: np.ndarray, params: TensorList, closure):
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return _ensure_float(closure(False))
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = self.get_params()
+
+        x0 = params.to_vec().detach().cpu().numpy()
+
+        res = directsearch.solve(
+            partial(self._objective, params = params, closure = closure),
+            x0 = x0,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return res.f
+
+
+
+class DirectSearchDS(Optimizer):
+    def __init__(
+        self,
+        params,
+        maxevals = DEFAULT_PARAMS['maxevals'], # Maximum number of function evaluations
+        rho = DEFAULT_PARAMS['rho'], # Forcing function
+        poll_type = DEFAULT_PARAMS['poll_type'], # Polling direction type
+        alpha0 = DEFAULT_PARAMS['alpha0'], # Original stepsize value
+        alpha_max = DEFAULT_PARAMS['alpha_max'], # Maximum value for the stepsize
+        alpha_min = DEFAULT_PARAMS['alpha_min'], # Minimum value for the stepsize
+        gamma_inc = DEFAULT_PARAMS['gamma_inc'], # Increasing factor for the stepsize
+        gamma_dec = DEFAULT_PARAMS['gamma_dec'], # Decreasing factor for the stepsize
+        verbose = DEFAULT_PARAMS['verbose'], # Display information about the method
+        print_freq = DEFAULT_PARAMS['print_freq'], # How frequently to display information
+        rho_uses_normd = DEFAULT_PARAMS['rho_uses_normd'], # Forcing function based on direction norm
+    ):
+        super().__init__(params, {})
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    def _objective(self, x: np.ndarray, params: TensorList, closure):
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return _ensure_float(closure(False))
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = self.get_params()
+
+        x0 = params.to_vec().detach().cpu().numpy()
+
+        res = directsearch.solve_directsearch(
+            partial(self._objective, params = params, closure = closure),
+            x0 = x0,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return res.f
+
+class DirectSearchProbabilistic(Optimizer):
+    def __init__(
+        self,
+        params,
+        maxevals = DEFAULT_PARAMS['maxevals'], # Maximum number of function evaluations
+        rho = DEFAULT_PARAMS['rho'], # Forcing function
+        alpha0 = DEFAULT_PARAMS['alpha0'], # Original stepsize value
+        alpha_max = DEFAULT_PARAMS['alpha_max'], # Maximum value for the stepsize
+        alpha_min = DEFAULT_PARAMS['alpha_min'], # Minimum value for the stepsize
+        gamma_inc = DEFAULT_PARAMS['gamma_inc'], # Increasing factor for the stepsize
+        gamma_dec = DEFAULT_PARAMS['gamma_dec'], # Decreasing factor for the stepsize
+        verbose = DEFAULT_PARAMS['verbose'], # Display information about the method
+        print_freq = DEFAULT_PARAMS['print_freq'], # How frequently to display information
+        rho_uses_normd = DEFAULT_PARAMS['rho_uses_normd'], # Forcing function based on direction norm
+    ):
+        super().__init__(params, {})
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    def _objective(self, x: np.ndarray, params: TensorList, closure):
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return _ensure_float(closure(False))
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = self.get_params()
+
+        x0 = params.to_vec().detach().cpu().numpy()
+
+        res = directsearch.solve_probabilistic_directsearch(
+            partial(self._objective, params = params, closure = closure),
+            x0 = x0,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return res.f
+
+
+class DirectSearchSubspace(Optimizer):
+    def __init__(
+        self,
+        params,
+        maxevals = DEFAULT_PARAMS['maxevals'], # Maximum number of function evaluations
+        rho = DEFAULT_PARAMS['rho'], # Forcing function
+        sketch_dim = DEFAULT_PARAMS['sketch_dim'], # Target dimension for sketching
+        sketch_type = DEFAULT_PARAMS['sketch_type'], # Sketching technique
+        poll_type = DEFAULT_PARAMS['poll_type'], # Polling direction type
+        alpha0 = DEFAULT_PARAMS['alpha0'], # Original stepsize value
+        alpha_max = DEFAULT_PARAMS['alpha_max'], # Maximum value for the stepsize
+        alpha_min = DEFAULT_PARAMS['alpha_min'], # Minimum value for the stepsize
+        gamma_inc = DEFAULT_PARAMS['gamma_inc'], # Increasing factor for the stepsize
+        gamma_dec = DEFAULT_PARAMS['gamma_dec'], # Decreasing factor for the stepsize
+        verbose = DEFAULT_PARAMS['verbose'], # Display information about the method
+        print_freq = DEFAULT_PARAMS['print_freq'], # How frequently to display information
+        rho_uses_normd = DEFAULT_PARAMS['rho_uses_normd'], # Forcing function based on direction norm
+    ):
+        super().__init__(params, {})
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    def _objective(self, x: np.ndarray, params: TensorList, closure):
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return _ensure_float(closure(False))
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = self.get_params()
+
+        x0 = params.to_vec().detach().cpu().numpy()
+
+        res = directsearch.solve_subspace_directsearch(
+            partial(self._objective, params = params, closure = closure),
+            x0 = x0,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return res.f
+
+
+
+class DirectSearchSTP(Optimizer):
+    def __init__(
+        self,
+        params,
+        maxevals = DEFAULT_PARAMS['maxevals'], # Maximum number of function evaluations
+        alpha0 = DEFAULT_PARAMS['alpha0'], # Original stepsize value
+        alpha_min = DEFAULT_PARAMS['alpha_min'], # Minimum value for the stepsize
+        verbose = DEFAULT_PARAMS['verbose'], # Display information about the method
+        print_freq = DEFAULT_PARAMS['print_freq'], # How frequently to display information
+    ):
+        super().__init__(params, {})
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    def _objective(self, x: np.ndarray, params: TensorList, closure):
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return _ensure_float(closure(False))
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = self.get_params()
+
+        x0 = params.to_vec().detach().cpu().numpy()
+
+        res = directsearch.solve_stp(
+            partial(self._objective, params = params, closure = closure),
+            x0 = x0,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return res.f