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

torchzero/modules/wrappers/optim_wrapper.py

@@ -10,34 +10,48 @@ class Wrap(Module):
     """
     Wraps a pytorch optimizer to use it as a module.
 
-    .. note::
-        Custom param groups are supported only by `set_param_groups`, settings passed to Modular will be ignored.
+    Note:
+        Custom param groups are supported only by ``set_param_groups``, settings passed to Modular will be applied to all parameters.
 
     Args:
         opt_fn (Callable[..., torch.optim.Optimizer] | torch.optim.Optimizer):
-            function that takes in parameters and returns the optimizer, for example :code:`torch.optim.Adam`
-            or :code:`lambda parameters: torch.optim.Adam(parameters, lr=1e-3)`
+            function that takes in parameters and returns the optimizer, for example ``torch.optim.Adam``
+            or ``lambda parameters: torch.optim.Adam(parameters, lr=1e-3)``
         *args:
         **kwargs:
-            Extra args to be passed to opt_fn. The function is called as :code:`opt_fn(parameters, *args, **kwargs)`.
+            Extra args to be passed to opt_fn. The function is called as ``opt_fn(parameters, *args, **kwargs)``.
+        use_param_groups:
+            Whether to pass settings passed to Modular to the wrapped optimizer.
 
-    Example:
-        wrapping pytorch_optimizer.StableAdamW
+            Note that settings to the first parameter are used for all parameters,
+            so if you specified per-parameter settings, they will be ignored.
 
-        .. code-block:: py
+    ### Example:
+    wrapping pytorch_optimizer.StableAdamW
 
-            from pytorch_optimizer import StableAdamW
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Wrap(StableAdamW, lr=1),
-                tz.m.Cautious(),
-                tz.m.LR(1e-2)
-            )
+    ```python
 
+    from pytorch_optimizer import StableAdamW
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Wrap(StableAdamW, lr=1),
+        tz.m.Cautious(),
+        tz.m.LR(1e-2)
+    )
+    ```
 
     """
-    def __init__(self, opt_fn: Callable[..., torch.optim.Optimizer] | torch.optim.Optimizer, *args, **kwargs):
-        super().__init__()
+
+    def __init__(
+        self,
+        opt_fn: Callable[..., torch.optim.Optimizer] | torch.optim.Optimizer,
+        *args,
+        use_param_groups: bool = True,
+        **kwargs,
+    ):
+        defaults = dict(use_param_groups=use_param_groups)
+        super().__init__(defaults=defaults)
+
         self._opt_fn = opt_fn
         self._opt_args = args
         self._opt_kwargs = kwargs
@@ -48,7 +62,7 @@ class Wrap(Module):
             self.optimizer = self._opt_fn
 
     def set_param_groups(self, param_groups):
-        self._custom_param_groups = param_groups
+        self._custom_param_groups = _make_param_groups(param_groups, differentiable=False)
         return super().set_param_groups(param_groups)
 
     @torch.no_grad
@@ -61,37 +75,29 @@ class Wrap(Module):
             param_groups = params if self._custom_param_groups is None else self._custom_param_groups
             self.optimizer = self._opt_fn(param_groups, *self._opt_args, **self._opt_kwargs)
 
+        # set optimizer per-parameter settings
+        if self.defaults["use_param_groups"] and var.modular is not None:
+            for group in self.optimizer.param_groups:
+                first_param = group['params'][0]
+                setting = self.settings[first_param]
+
+                # settings passed in `set_param_groups` are the highest priority
+                # schedulers will override defaults but not settings passed in `set_param_groups`
+                # this is consistent with how Modular does it.
+                if self._custom_param_groups is not None:
+                    setting = {k:v for k,v in setting if k not in self._custom_param_groups[0]}
+
+                group.update(setting)
+
         # set grad to update
         orig_grad = [p.grad for p in params]
         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 var.is_last and (
-            (var.last_module_lrs is None)
-            or
-            (('lr' in self.optimizer.defaults) and (len(set(var.last_module_lrs)) == 1))
-        ):
-            lr = 1 if var.last_module_lrs is None else var.last_module_lrs[0]
-
-            # update optimizer lr with desired lr
-            if lr != 1:
-                self.optimizer.defaults['__original_lr__'] = self.optimizer.defaults['lr']
-                for g in self.optimizer.param_groups:
-                    g['__original_lr__'] = g['lr']
-                    g['lr'] = g['lr'] * lr
-
-            # step
+        # if this is last module, simply use optimizer to update parameters
+        if var.modular is not None and self is var.modular.modules[-1]:
             self.optimizer.step()
 
-            # restore original lr
-            if lr != 1:
-                self.optimizer.defaults['lr'] = self.optimizer.defaults.pop('__original_lr__')
-                for g in self.optimizer.param_groups:
-                    g['lr'] = g.pop('__original_lr__')
-
             # restore grad
             for p, g in zip(params, orig_grad):
                 p.grad = g
@@ -100,6 +106,7 @@ class Wrap(Module):
             return var
 
         # this is not the last module, meaning update is difference in parameters
+        # and passed to next module
         params_before_step = [p.clone() for p in params]
         self.optimizer.step() # step and update params
         for p, g in zip(params, orig_grad):

torchzero/modules/zeroth_order/__init__.py

@@ -1 +1 @@
-from .cd import CD, CCD, CCDLS
+from .cd import CD

torchzero/modules/zeroth_order/cd.py

@@ -9,11 +9,10 @@ import torch
 
 from ...core import Module
 from ...utils import NumberList, TensorList
-from ..line_search.adaptive import adaptive_tracking
 
 class CD(Module):
     """Coordinate descent. Proposes a descent direction along a single coordinate.
-    You can then put a line search such as ``tz.m.ScipyMinimizeScalar``, or a fixed step size.
+    A line search such as ``tz.m.ScipyMinimizeScalar(maxiter=8)`` or a fixed step size can be used after this.
 
     Args:
         h (float, optional): finite difference step size. Defaults to 1e-3.
@@ -121,239 +120,3 @@ class CD(Module):
         var.update = update
         return var
 
-
-def _icd_get_idx(self: Module, params: TensorList):
-    ndim = params.global_numel()
-    igrad = self.get_state(params, "igrad", cls=TensorList)
-
-    # -------------------------- 1st n steps fill igrad -------------------------- #
-    index = self.global_state.get('index', 0)
-    self.global_state['index'] = index + 1
-    if index < ndim:
-        return index, igrad
-
-    # ------------------ select randomly weighted by magnitudes ------------------ #
-    igrad_abs = igrad.abs()
-    gmin = igrad_abs.global_min()
-    gmax = igrad_abs.global_max()
-
-    pmin, pmax, pow = self.get_settings(params, "pmin", "pmax", "pow", cls=NumberList)
-
-    p: TensorList = ((igrad_abs - gmin) / (gmax - gmin)) ** pow # pyright:ignore[reportOperatorIssue]
-    p.mul_(pmax-pmin).add_(pmin)
-
-    if 'np_gen' not in self.global_state:
-        self.global_state['np_gen'] = np.random.default_rng(0)
-    np_gen = self.global_state['np_gen']
-
-    p_vec = p.to_vec()
-    p_sum = p_vec.sum()
-    if p_sum > 1e-12:
-        return np_gen.choice(ndim, p=p_vec.div_(p_sum).numpy(force=True)), igrad
-
-    # --------------------- sum is too small, do cycle again --------------------- #
-    self.global_state.clear()
-    self.clear_state_keys('h_vec', 'igrad', 'alphas')
-
-    if 'generator' not in self.global_state:
-        self.global_state['generator'] = random.Random(0)
-    generator = self.global_state['generator']
-    return generator.randrange(0, p_vec.numel()), igrad
-
-class CCD(Module):
-    """Cumulative coordinate descent. This updates one gradient coordinate at a time and accumulates it
-    to the update direction. The coordinate updated is random weighted by magnitudes of current update direction.
-    As update direction ceases to be a descent direction due to old accumulated coordinates, it is decayed.
-
-    Args:
-        pmin (float, optional): multiplier to probability of picking the lowest magnitude gradient. Defaults to 0.1.
-        pmax (float, optional): multiplier to probability of picking the largest magnitude gradient. Defaults to 1.0.
-        pow (int, optional): power transform to probabilities. Defaults to 2.
-        decay (float, optional): accumulated gradient decay on failed step. Defaults to 0.5.
-        decay2 (float, optional): decay multiplier decay on failed step. Defaults to 0.25.
-        nplus (float, optional): step size increase on successful steps. Defaults to 1.5.
-        nminus (float, optional): step size increase on unsuccessful steps. Defaults to 0.75.
-    """
-    def __init__(self, pmin=0.1, pmax=1.0, pow=2, decay:float=0.8, decay2:float=0.2, nplus=1.5, nminus=0.75):
-
-        defaults = dict(pmin=pmin, pmax=pmax, pow=pow, decay=decay, decay2=decay2, nplus=nplus, nminus=nminus)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        if closure is None:
-            raise RuntimeError("CD requires closure")
-
-        params = TensorList(var.params)
-        p_prev = self.get_state(params, "p_prev", init=params, cls=TensorList)
-
-        f_0 = var.get_loss(False)
-        step_size = self.global_state.get('step_size', 1)
-
-        # ------------------------ hard reset on infinite loss ----------------------- #
-        if not math.isfinite(f_0):
-            del self.global_state['f_prev']
-            var.update = params - p_prev
-            self.global_state.clear()
-            self.state.clear()
-            self.global_state["step_size"] = step_size / 10
-            return var
-
-        # ---------------------------- soft reset if stuck --------------------------- #
-        if "igrad" in self.state[params[0]]:
-            n_bad = self.global_state.get('n_bad', 0)
-
-            f_prev = self.global_state.get("f_prev", None)
-            if f_prev is not None:
-
-                decay2 = self.defaults["decay2"]
-                decay = self.global_state.get("decay", self.defaults["decay"])
-
-                if f_0 >= f_prev:
-
-                    igrad = self.get_state(params, "igrad", cls=TensorList)
-                    del self.global_state['f_prev']
-
-                    # undo previous update
-                    var.update = params - p_prev
-
-                    # increment n_bad
-                    self.global_state['n_bad'] = n_bad + 1
-
-                    # decay step size
-                    self.global_state['step_size'] = step_size * self.defaults["nminus"]
-
-                    # soft reset
-                    if n_bad > 0:
-                        igrad *= decay
-                        self.global_state["decay"] = decay*decay2
-                        self.global_state['n_bad'] = 0
-
-                    return var
-
-                else:
-                    # increase step size and reset n_bad
-                    self.global_state['step_size'] = step_size * self.defaults["nplus"]
-                    self.global_state['n_bad'] = 0
-                    self.global_state["decay"] = self.defaults["decay"]
-
-            self.global_state['f_prev'] = float(f_0)
-
-        # ------------------------------ determine index ----------------------------- #
-        idx, igrad = _icd_get_idx(self, params)
-
-        # -------------------------- find descent direction -------------------------- #
-        h_vec = self.get_state(params, 'h_vec', init=lambda x: torch.full_like(x, 1e-3), cls=TensorList)
-        h = float(h_vec.flat_get(idx))
-
-        params.flat_set_lambda_(idx, lambda x: x + h)
-        f_p = closure(False)
-
-        params.flat_set_lambda_(idx, lambda x: x - 2*h)
-        f_n = closure(False)
-        params.flat_set_lambda_(idx, lambda x: x + h)
-
-        # ---------------------------------- adapt h --------------------------------- #
-        if f_0 <= f_p and f_0 <= f_n:
-            h_vec.flat_set_lambda_(idx, lambda x: max(x/2, 1e-10))
-        else:
-            if abs(f_0 - f_n) < 1e-12 or abs((f_p - f_0) / (f_0 - f_n) - 1) < 1e-2:
-                h_vec.flat_set_lambda_(idx, lambda x: min(x*2, 1e10))
-
-        # ------------------------------- update igrad ------------------------------- #
-        if f_0 < f_p and f_0 < f_n: alpha = 0
-        else: alpha = (f_p - f_n) / (2*h)
-
-        igrad.flat_set_(idx, alpha)
-
-        # ----------------------------- create the update ---------------------------- #
-        var.update = igrad * step_size
-        p_prev.copy_(params)
-        return var
-
-
-class CCDLS(Module):
-    """CCD with line search instead of adaptive step size.
-
-    Args:
-        pmin (float, optional): multiplier to probability of picking the lowest magnitude gradient. Defaults to 0.1.
-        pmax (float, optional): multiplier to probability of picking the largest magnitude gradient. Defaults to 1.0.
-        pow (int, optional): power transform to probabilities. Defaults to 2.
-        decay (float, optional): accumulated gradient decay on failed step. Defaults to 0.5.
-        decay2 (float, optional): decay multiplier decay on failed step. Defaults to 0.25.
-        maxiter (int, optional): max number of line search iterations.
-    """
-    def __init__(self, pmin=0.1, pmax=1.0, pow=2, decay=0.8, decay2=0.2, maxiter=10, ):
-        defaults = dict(pmin=pmin, pmax=pmax, pow=pow, maxiter=maxiter, decay=decay, decay2=decay2)
-        super().__init__(defaults)
-
-    @torch.no_grad
-    def step(self, var):
-        closure = var.closure
-        if closure is None:
-            raise RuntimeError("CD requires closure")
-
-        params = TensorList(var.params)
-        finfo = torch.finfo(params[0].dtype)
-        f_0 = var.get_loss(False)
-
-        # ------------------------------ determine index ----------------------------- #
-        idx, igrad = _icd_get_idx(self, params)
-
-        # -------------------------- find descent direction -------------------------- #
-        h_vec = self.get_state(params, 'h_vec', init=lambda x: torch.full_like(x, 1e-3), cls=TensorList)
-        h = float(h_vec.flat_get(idx))
-
-        params.flat_set_lambda_(idx, lambda x: x + h)
-        f_p = closure(False)
-
-        params.flat_set_lambda_(idx, lambda x: x - 2*h)
-        f_n = closure(False)
-        params.flat_set_lambda_(idx, lambda x: x + h)
-
-        # ---------------------------------- adapt h --------------------------------- #
-        if f_0 <= f_p and f_0 <= f_n:
-            h_vec.flat_set_lambda_(idx, lambda x: max(x/2, finfo.tiny * 2))
-        else:
-            # here eps, not tiny
-            if abs(f_0 - f_n) < finfo.eps or abs((f_p - f_0) / (f_0 - f_n) - 1) < 1e-2:
-                h_vec.flat_set_lambda_(idx, lambda x: min(x*2, finfo.max / 2))
-
-        # ------------------------------- update igrad ------------------------------- #
-        if f_0 < f_p and f_0 < f_n: alpha = 0
-        else: alpha = (f_p - f_n) / (2*h)
-
-        igrad.flat_set_(idx, alpha)
-
-        # -------------------------------- line search ------------------------------- #
-        x0 = params.clone()
-        def f(a):
-            params.sub_(igrad, alpha=a)
-            loss = closure(False)
-            params.copy_(x0)
-            return loss
-
-        a_prev = self.global_state.get('a_prev', 1)
-        a, f_a, niter = adaptive_tracking(f, a_prev, maxiter=self.defaults['maxiter'], f_0=f_0)
-        if (a is None) or (not math.isfinite(a)) or (not math.isfinite(f_a)):
-            a = 0
-
-        # -------------------------------- set a_prev -------------------------------- #
-        decay2 = self.defaults["decay2"]
-        decay = self.global_state.get("decay", self.defaults["decay"])
-
-        if abs(a) > finfo.tiny * 2:
-            assert f_a < f_0
-            self.global_state['a_prev'] = max(min(a, finfo.max / 2), finfo.tiny * 2)
-            self.global_state["decay"] = self.defaults["decay"]
-
-        # ---------------------------- soft reset on fail ---------------------------- #
-        else:
-            igrad *= decay
-            self.global_state["decay"] = decay*decay2
-            self.global_state['a_prev'] = a_prev / 2
-
-        # -------------------------------- set update -------------------------------- #
-        var.update = igrad * a
-        return var

torchzero/utils/derivatives.py

@@ -5,13 +5,13 @@ import torch.autograd.forward_ad as fwAD
 
 from .torch_tools import swap_tensors_no_use_count_check, vec_to_tensors
 
-def _jacobian(output: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False):
-    flat_input = torch.cat([i.reshape(-1) for i in output])
-    grad_ouputs = torch.eye(len(flat_input), device=output[0].device, dtype=output[0].dtype)
+def _jacobian(outputs: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False):
+    flat_outputs = torch.cat([i.reshape(-1) for i in outputs])
+    grad_ouputs = torch.eye(len(flat_outputs), device=outputs[0].device, dtype=outputs[0].dtype)
     jac = []
-    for i in range(flat_input.numel()):
+    for i in range(flat_outputs.numel()):
         jac.append(torch.autograd.grad(
-            flat_input,
+            flat_outputs,
             wrt,
             grad_ouputs[i],
             retain_graph=True,
@@ -22,12 +22,12 @@ def _jacobian(output: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], creat
     return [torch.stack(z) for z in zip(*jac)]
 
 
-def _jacobian_batched(output: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False):
-    flat_input = torch.cat([i.reshape(-1) for i in output])
+def _jacobian_batched(outputs: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False):
+    flat_outputs = torch.cat([i.reshape(-1) for i in outputs])
     return torch.autograd.grad(
-        flat_input,
+        flat_outputs,
         wrt,
-        torch.eye(len(flat_input), device=output[0].device, dtype=output[0].dtype),
+        torch.eye(len(flat_outputs), device=outputs[0].device, dtype=outputs[0].dtype),
         retain_graph=True,
         create_graph=create_graph,
         allow_unused=True,
@@ -51,13 +51,13 @@ def flatten_jacobian(jacs: Sequence[torch.Tensor]) -> torch.Tensor:
     return torch.cat([j.reshape(n_out, -1) for j in jacs], dim=1)
 
 
-def jacobian_wrt(output: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False, batched=True) -> Sequence[torch.Tensor]:
+def jacobian_wrt(outputs: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False, batched=True) -> Sequence[torch.Tensor]:
     """Calculate jacobian of a sequence of tensors w.r.t another sequence of tensors.
     Returns a sequence of tensors with the length as `wrt`.
     Each tensor will have the shape `(*output.shape, *wrt[i].shape)`.
 
     Args:
-        input (Sequence[torch.Tensor]): input sequence of tensors.
+        outputs (Sequence[torch.Tensor]): input sequence of tensors.
         wrt (Sequence[torch.Tensor]): sequence of tensors to differentiate w.r.t.
         create_graph (bool, optional):
             pytorch option, if True, graph of the derivative will be constructed,
@@ -68,16 +68,16 @@ def jacobian_wrt(output: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], cr
     Returns:
         sequence of tensors with the length as `wrt`.
     """
-    if batched: return _jacobian_batched(output, wrt, create_graph)
-    return _jacobian(output, wrt, create_graph)
+    if batched: return _jacobian_batched(outputs, wrt, create_graph)
+    return _jacobian(outputs, wrt, create_graph)
 
-def jacobian_and_hessian_wrt(output: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False, batched=True):
+def jacobian_and_hessian_wrt(outputs: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False, batched=True):
     """Calculate jacobian and hessian of a sequence of tensors w.r.t another sequence of tensors.
     Calculating hessian requires calculating the jacobian. So this function is more efficient than
     calling `jacobian` and `hessian` separately, which would calculate jacobian twice.
 
     Args:
-        input (Sequence[torch.Tensor]): input sequence of tensors.
+        outputs (Sequence[torch.Tensor]): input sequence of tensors.
         wrt (Sequence[torch.Tensor]): sequence of tensors to differentiate w.r.t.
         create_graph (bool, optional):
             pytorch option, if True, graph of the derivative will be constructed,
@@ -87,7 +87,7 @@ def jacobian_and_hessian_wrt(output: Sequence[torch.Tensor], wrt: Sequence[torch
     Returns:
         tuple with jacobians sequence and hessians sequence.
     """
-    jac = jacobian_wrt(output, wrt, create_graph=True, batched = batched)
+    jac = jacobian_wrt(outputs, wrt, create_graph=True, batched = batched)
     return jac, jacobian_wrt(jac, wrt, batched = batched, create_graph=create_graph)
 
 
@@ -96,13 +96,13 @@ def jacobian_and_hessian_wrt(output: Sequence[torch.Tensor], wrt: Sequence[torch
 #     Note - I only tested this for cases where input is a scalar."""
 #     return torch.cat([h.reshape(h.size(0), h[1].numel()) for h in hessians], 1)
 
-def jacobian_and_hessian_mat_wrt(output: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False, batched=True):
+def jacobian_and_hessian_mat_wrt(outputs: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor], create_graph=False, batched=True):
     """Calculate jacobian and hessian of a sequence of tensors w.r.t another sequence of tensors.
     Calculating hessian requires calculating the jacobian. So this function is more efficient than
     calling `jacobian` and `hessian` separately, which would calculate jacobian twice.
 
     Args:
-        input (Sequence[torch.Tensor]): input sequence of tensors.
+        outputs (Sequence[torch.Tensor]): input sequence of tensors.
         wrt (Sequence[torch.Tensor]): sequence of tensors to differentiate w.r.t.
         create_graph (bool, optional):
             pytorch option, if True, graph of the derivative will be constructed,
@@ -112,7 +112,7 @@ def jacobian_and_hessian_mat_wrt(output: Sequence[torch.Tensor], wrt: Sequence[t
     Returns:
         tuple with jacobians sequence and hessians sequence.
     """
-    jac = jacobian_wrt(output, wrt, create_graph=True, batched = batched)
+    jac = jacobian_wrt(outputs, wrt, create_graph=True, batched = batched)
     H_list = jacobian_wrt(jac, wrt, batched = batched, create_graph=create_graph)
     return flatten_jacobian(jac), flatten_jacobian(H_list)
 

torchzero/utils/linalg/linear_operator.py

@@ -35,8 +35,8 @@ class LinearOperator(ABC):
         """solve with a norm bound on x"""
         raise NotImplementedError(f"{self.__class__.__name__} doesn't implement solve_bounded")
 
-    def update(self, *args, **kwargs) -> None:
-        raise NotImplementedError(f"{self.__class__.__name__} doesn't implement update")
+    # def update(self, *args, **kwargs) -> None:
+    #     raise NotImplementedError(f"{self.__class__.__name__} doesn't implement update")
 
     def add(self, x: torch.Tensor) -> "LinearOperator":
         raise NotImplementedError(f"{self.__class__.__name__} doesn't implement add")
@@ -298,6 +298,7 @@ class AtA(LinearOperator):
 class AAT(LinearOperator):
     def __init__(self, A: torch.Tensor):
         self.A = A
+        self.device = self.A.device; self.dtype = self.A.dtype
 
     def matvec(self, x): return self.A.mv(self.A.mH.mv(x))
     def rmatvec(self, x): return self.matvec(x)
@@ -327,3 +328,50 @@ class AAT(LinearOperator):
         n = self.A.size(1)
         return (n,n)
 
+
+class Sketched(LinearOperator):
+    """A projected by sketching matrix S, representing the operator S @ A_proj @ S.T.
+
+    Where A is (n, n) and S is (n, sketch_size).
+    """
+    def __init__(self, S: torch.Tensor, A_proj: torch.Tensor):
+        self.S = S
+        self.A_proj = A_proj
+        self.device = self.A_proj.device; self.dtype = self.A_proj.dtype
+
+
+    def matvec(self, x):
+        x_proj = self.S.T @ x
+        Ax_proj = self.A_proj @ x_proj
+        return self.S @ Ax_proj
+
+    def rmatvec(self, x):
+        x_proj = self.S.T @ x
+        ATx_proj = self.A_proj.mH @ x_proj
+        return self.S @ ATx_proj
+
+
+    def matmat(self, x): return Dense(torch.linalg.multi_dot([self.S, self.A_proj, self.S.T, x])) # pylint:disable=not-callable
+    def rmatmat(self, x): return Dense(torch.linalg.multi_dot([self.S, self.A_proj.mH, self.S.T, x])) # pylint:disable=not-callable
+
+
+    def is_dense(self): return False
+    def to_tensor(self): return self.S @ self.A_proj @ self.S.T
+    def transpose(self): return Sketched(self.S, self.A_proj.mH)
+
+    def add_diagonal(self, x):
+        """this doesn't correspond to adding diagonal to A, however it still works for LM etc."""
+        if isinstance(x, torch.Tensor) and x.numel() <= 1: x = x.item()
+        if isinstance(x, (int,float)): x = torch.full((self.A_proj.shape[0],), fill_value=x, device=self.A_proj.device, dtype=self.A_proj.dtype)
+        return Sketched(S=self.S, A_proj=self.A_proj + x.diag_embed())
+
+    def solve(self, b):
+        return self.S @ torch.linalg.lstsq(self.A_proj, self.S.T @ b).solution # pylint:disable=not-callable
+
+    def inv(self):
+        return Sketched(S=self.S, A_proj=torch.linalg.pinv(self.A_proj)) # pylint:disable=not-callable
+
+    def size(self):
+        n = self.S.size(0)
+        return (n,n)
+

torchzero/utils/optimizer.py

@@ -110,7 +110,7 @@ def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], param
         for i, param in enumerate(params):
             s = state[param]
             if key not in s:
-                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                if must_exist: raise KeyError(f"Key `{key}` doesn't exist in state with keys {tuple(s.keys())}")
                 s[key] = _make_initial_state_value(param, init, i)
             values.append(s[key])
         return values
@@ -125,7 +125,7 @@ def get_state_vals(state: Mapping[torch.Tensor, MutableMapping[str, Any]], param
         s = state[param]
         for k_i, key in enumerate(keys):
             if key not in s:
-                if must_exist: raise KeyError(f"Key {key} doesn't exist in state with keys {tuple(s.keys())}")
+                if must_exist: raise KeyError(f"Key `{key}` doesn't exist in state with keys {tuple(s.keys())}")
                 k_init = init[k_i] if isinstance(init, (list,tuple)) else init
                 s[key] = _make_initial_state_value(param, k_init, i)
             values[k_i].append(s[key])

torchzero/utils/python_tools.py

@@ -67,3 +67,4 @@ def safe_dict_update_(d1_:dict, d2:dict):
     inter = set(d1_.keys()).intersection(d2.keys())
     if len(inter) > 0: raise RuntimeError(f"Duplicate keys {inter}")
     d1_.update(d2)
+

{torchzero-0.3.13.dist-info → torchzero-0.3.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchzero
-Version: 0.3.13
+Version: 0.3.15
 Summary: Modular optimization library for PyTorch.
 Author-email: Ivan Nikishev <nkshv2@gmail.com>
 Project-URL: Homepage, https://github.com/inikishev/torchzero