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

torchzero/utils/__init__.py

@@ -1,5 +1,11 @@
 from . import tensorlist as tl
-from .compile import _optional_compiler, benchmark_compile_cpu, benchmark_compile_cuda, set_compilation, enable_compilation
+from .compile import (
+    _optional_compiler,
+    benchmark_compile_cpu,
+    benchmark_compile_cuda,
+    enable_compilation,
+    set_compilation,
+)
 from .numberlist import NumberList
 from .optimizer import (
     Init,
@@ -18,6 +24,36 @@ from .params import (
     _copy_param_groups,
     _make_param_groups,
 )
-from .python_tools import flatten, generic_eq, generic_ne, reduce_dim, unpack_dicts
-from .tensorlist import TensorList, as_tensorlist, Distributions, generic_clamp, generic_numel, generic_vector_norm, generic_zeros_like, generic_randn_like, generic_finfo_eps
-from .torch_tools import tofloat, tolist, tonumpy, totensor, vec_to_tensors, vec_to_tensors_, set_storage_
+from .python_tools import (
+    flatten,
+    generic_eq,
+    generic_ne,
+    reduce_dim,
+    safe_dict_update_,
+    unpack_dicts,
+)
+from .tensorlist import (
+    Distributions,
+    Metrics,
+    TensorList,
+    as_tensorlist,
+    generic_clamp,
+    generic_finfo,
+    generic_finfo_eps,
+    generic_finfo_tiny,
+    generic_max,
+    generic_numel,
+    generic_randn_like,
+    generic_sum,
+    generic_vector_norm,
+    generic_zeros_like,
+)
+from .torch_tools import (
+    set_storage_,
+    tofloat,
+    tolist,
+    tonumpy,
+    totensor,
+    vec_to_tensors,
+    vec_to_tensors_,
+)

torchzero/utils/compile.py

@@ -38,7 +38,7 @@ class _MaybeCompiledFunc:
 _optional_compiler = _OptionalCompiler()
 """this holds .enable attribute, set to True to enable compiling for a few functions that benefit from it."""
 
-def set_compilation(enable: bool):
+def set_compilation(enable: bool=True):
     """`enable` is False by default. When True, certain functions will be compiled, which may not work on some systems like Windows, but it usually improves performance."""
     _optional_compiler.enable = enable
 

torchzero/utils/derivatives.py

@@ -2,7 +2,6 @@ from collections.abc import Iterable, Sequence
 
 import torch
 import torch.autograd.forward_ad as fwAD
-from typing import Literal
 
 from .torch_tools import swap_tensors_no_use_count_check, vec_to_tensors
 
@@ -35,10 +34,27 @@ def _jacobian_batched(output: Sequence[torch.Tensor], wrt: Sequence[torch.Tensor
         is_grads_batched=True,
     )
 
+def flatten_jacobian(jacs: Sequence[torch.Tensor]) -> torch.Tensor:
+    """Converts the output of jacobian_wrt (a list of tensors) into a single 2D matrix.
+
+    Args:
+        jacs (Sequence[torch.Tensor]):
+            output from jacobian_wrt where ach tensor has the shape `(*output.shape, *wrt[i].shape)`.
+
+    Returns:
+        torch.Tensor: has the shape `(output.ndim, wrt.ndim)`.
+    """
+    if not jacs:
+        return torch.empty(0, 0)
+
+    n_out = jacs[0].shape[0]
+    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]:
     """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 `(*input.shape, *wrt[i].shape)`.
+    Each tensor will have the shape `(*output.shape, *wrt[i].shape)`.
 
     Args:
         input (Sequence[torch.Tensor]): input sequence of tensors.
@@ -75,10 +91,10 @@ def jacobian_and_hessian_wrt(output: Sequence[torch.Tensor], wrt: Sequence[torch
     return jac, jacobian_wrt(jac, wrt, batched = batched, create_graph=create_graph)
 
 
-def hessian_list_to_mat(hessians: Sequence[torch.Tensor]):
-    """takes output of `hessian` and returns the 2D hessian matrix.
-    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 hessian_list_to_mat(hessians: Sequence[torch.Tensor]):
+#     """takes output of `hessian` and returns the 2D hessian matrix.
+#     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):
     """Calculate jacobian and hessian of a sequence of tensors w.r.t another sequence of tensors.
@@ -98,7 +114,7 @@ def jacobian_and_hessian_mat_wrt(output: Sequence[torch.Tensor], wrt: Sequence[t
     """
     jac = jacobian_wrt(output, wrt, create_graph=True, batched = batched)
     H_list = jacobian_wrt(jac, wrt, batched = batched, create_graph=create_graph)
-    return torch.cat([j.view(-1) for j in jac]), hessian_list_to_mat(H_list)
+    return flatten_jacobian(jac), flatten_jacobian(H_list)
 
 def hessian(
     fn,
@@ -115,19 +131,18 @@ def hessian(
     `vectorize` and `outer_jacobian_strategy` are only for `method = "torch.autograd"`, refer to its documentation.
 
     Example:
-    .. code:: py
-
-        model = nn.Linear(4, 2) # (2, 4) weight and (2, ) bias
-        X = torch.randn(10, 4)
-        y = torch.randn(10, 2)
+    ```python
+    model = nn.Linear(4, 2) # (2, 4) weight and (2, ) bias
+    X = torch.randn(10, 4)
+    y = torch.randn(10, 2)
 
-        def fn():
-            y_hat = model(X)
-            loss = F.mse_loss(y_hat, y)
-            return loss
-
-        hessian_mat(fn, model.parameters()) # list of two lists of two lists of 3D and 4D tensors
+    def fn():
+        y_hat = model(X)
+        loss = F.mse_loss(y_hat, y)
+        return loss
 
+    hessian_mat(fn, model.parameters()) # list of two lists of two lists of 3D and 4D tensors
+    ```
 
     """
     params = list(params)
@@ -165,19 +180,18 @@ def hessian_mat(
     `vectorize` and `outer_jacobian_strategy` are only for `method = "torch.autograd"`, refer to its documentation.
 
     Example:
-    .. code:: py
-
-        model = nn.Linear(4, 2) # 10 parameters in total
-        X = torch.randn(10, 4)
-        y = torch.randn(10, 2)
+    ```python
+    model = nn.Linear(4, 2) # 10 parameters in total
+    X = torch.randn(10, 4)
+    y = torch.randn(10, 2)
 
-        def fn():
-            y_hat = model(X)
-            loss = F.mse_loss(y_hat, y)
-            return loss
-
-        hessian_mat(fn, model.parameters()) # 10x10 tensor
+    def fn():
+        y_hat = model(X)
+        loss = F.mse_loss(y_hat, y)
+        return loss
 
+    hessian_mat(fn, model.parameters()) # 10x10 tensor
+    ```
 
     """
     params = list(params)
@@ -206,21 +220,20 @@ def jvp(fn, params: Iterable[torch.Tensor], tangent: Iterable[torch.Tensor]) ->
     """Jacobian vector product.
 
     Example:
-    .. code:: py
-
-        model = nn.Linear(4, 2)
-        X = torch.randn(10, 4)
-        y = torch.randn(10, 2)
-
-        tangent = [torch.randn_like(p) for p in model.parameters()]
+    ```python
+    model = nn.Linear(4, 2)
+    X = torch.randn(10, 4)
+    y = torch.randn(10, 2)
 
-        def fn():
-            y_hat = model(X)
-            loss = F.mse_loss(y_hat, y)
-            return loss
+    tangent = [torch.randn_like(p) for p in model.parameters()]
 
-        jvp(fn, model.parameters(), tangent) # scalar
+    def fn():
+        y_hat = model(X)
+        loss = F.mse_loss(y_hat, y)
+        return loss
 
+    jvp(fn, model.parameters(), tangent) # scalar
+    ```
     """
     params = list(params)
     tangent = list(tangent)
@@ -253,21 +266,20 @@ def jvp_fd_central(
     """Jacobian vector product using central finite difference formula.
 
     Example:
-    .. code:: py
-
-        model = nn.Linear(4, 2)
-        X = torch.randn(10, 4)
-        y = torch.randn(10, 2)
-
-        tangent = [torch.randn_like(p) for p in model.parameters()]
+    ```python
+    model = nn.Linear(4, 2)
+    X = torch.randn(10, 4)
+    y = torch.randn(10, 2)
 
-        def fn():
-            y_hat = model(X)
-            loss = F.mse_loss(y_hat, y)
-            return loss
+    tangent = [torch.randn_like(p) for p in model.parameters()]
 
-        jvp_fd_central(fn, model.parameters(), tangent) # scalar
+    def fn():
+        y_hat = model(X)
+        loss = F.mse_loss(y_hat, y)
+        return loss
 
+    jvp_fd_central(fn, model.parameters(), tangent) # scalar
+    ```
     """
     params = list(params)
     tangent = list(tangent)
@@ -304,24 +316,24 @@ def jvp_fd_forward(
     Loss at initial point can be specified in the `v_0` argument.
 
     Example:
-    .. code:: py
+    ```python
+    model = nn.Linear(4, 2)
+    X = torch.randn(10, 4)
+    y = torch.randn(10, 2)
 
-        model = nn.Linear(4, 2)
-        X = torch.randn(10, 4)
-        y = torch.randn(10, 2)
+    tangent1 = [torch.randn_like(p) for p in model.parameters()]
+    tangent2 = [torch.randn_like(p) for p in model.parameters()]
 
-        tangent1 = [torch.randn_like(p) for p in model.parameters()]
-        tangent2 = [torch.randn_like(p) for p in model.parameters()]
-
-        def fn():
-            y_hat = model(X)
-            loss = F.mse_loss(y_hat, y)
-            return loss
+    def fn():
+        y_hat = model(X)
+        loss = F.mse_loss(y_hat, y)
+        return loss
 
-        v_0 = fn() # pre-calculate loss at initial point
+    v_0 = fn() # pre-calculate loss at initial point
 
-        jvp1 = jvp_fd_forward(fn, model.parameters(), tangent1, v_0=v_0) # scalar
-        jvp2 = jvp_fd_forward(fn, model.parameters(), tangent2, v_0=v_0) # scalar
+    jvp1 = jvp_fd_forward(fn, model.parameters(), tangent1, v_0=v_0) # scalar
+    jvp2 = jvp_fd_forward(fn, model.parameters(), tangent2, v_0=v_0) # scalar
+    ```
 
     """
     params = list(params)
@@ -356,21 +368,21 @@ def hvp(
     """Hessian-vector product
 
     Example:
-    .. code:: py
-
-        model = nn.Linear(4, 2)
-        X = torch.randn(10, 4)
-        y = torch.randn(10, 2)
+    ```python
+    model = nn.Linear(4, 2)
+    X = torch.randn(10, 4)
+    y = torch.randn(10, 2)
 
-        y_hat = model(X)
-        loss = F.mse_loss(y_hat, y)
-        loss.backward(create_graph=True)
+    y_hat = model(X)
+    loss = F.mse_loss(y_hat, y)
+    loss.backward(create_graph=True)
 
-        grads = [p.grad for p in model.parameters()]
-        vec = [torch.randn_like(p) for p in model.parameters()]
+    grads = [p.grad for p in model.parameters()]
+    vec = [torch.randn_like(p) for p in model.parameters()]
 
-        # list of tensors, same layout as model.parameters()
-        hvp(model.parameters(), grads, vec=vec)
+    # list of tensors, same layout as model.parameters()
+    hvp(model.parameters(), grads, vec=vec)
+    ```
     """
     params = list(params)
     g = list(grads)
@@ -393,23 +405,23 @@ def hvp_fd_central(
     Please note that this will clear :code:`grad` attributes in params.
 
     Example:
-    .. code:: py
-
-        model = nn.Linear(4, 2)
-        X = torch.randn(10, 4)
-        y = torch.randn(10, 2)
+    ```python
+    model = nn.Linear(4, 2)
+    X = torch.randn(10, 4)
+    y = torch.randn(10, 2)
 
-        def closure():
-            y_hat = model(X)
-            loss = F.mse_loss(y_hat, y)
-            model.zero_grad()
-            loss.backward()
-            return loss
+    def closure():
+        y_hat = model(X)
+        loss = F.mse_loss(y_hat, y)
+        model.zero_grad()
+        loss.backward()
+        return loss
 
-        vec = [torch.randn_like(p) for p in model.parameters()]
+    vec = [torch.randn_like(p) for p in model.parameters()]
 
-        # list of tensors, same layout as model.parameters()
-        hvp_fd_central(closure, model.parameters(), vec=vec)
+    # list of tensors, same layout as model.parameters()
+    hvp_fd_central(closure, model.parameters(), vec=vec)
+    ```
     """
     params = list(params)
     vec = list(vec)
@@ -456,27 +468,27 @@ def hvp_fd_forward(
     Please note that this will clear :code:`grad` attributes in params.
 
     Example:
-    .. code:: py
+    ```python
+    model = nn.Linear(4, 2)
+    X = torch.randn(10, 4)
+    y = torch.randn(10, 2)
 
-        model = nn.Linear(4, 2)
-        X = torch.randn(10, 4)
-        y = torch.randn(10, 2)
-
-        def closure():
-            y_hat = model(X)
-            loss = F.mse_loss(y_hat, y)
-            model.zero_grad()
-            loss.backward()
-            return loss
+    def closure():
+        y_hat = model(X)
+        loss = F.mse_loss(y_hat, y)
+        model.zero_grad()
+        loss.backward()
+        return loss
 
-        vec = [torch.randn_like(p) for p in model.parameters()]
+    vec = [torch.randn_like(p) for p in model.parameters()]
 
-        # pre-compute gradient at initial point
-        closure()
-        g_0 = [p.grad for p in model.parameters()]
+    # pre-compute gradient at initial point
+    closure()
+    g_0 = [p.grad for p in model.parameters()]
 
-        # list of tensors, same layout as model.parameters()
-        hvp_fd_forward(closure, model.parameters(), vec=vec, g_0=g_0)
+    # list of tensors, same layout as model.parameters()
+    hvp_fd_forward(closure, model.parameters(), vec=vec, g_0=g_0)
+    ```
     """
 
     params = list(params)
@@ -485,7 +497,7 @@ def hvp_fd_forward(
 
     vec_norm = None
     if normalize:
-        vec_norm = torch.linalg.vector_norm(torch.cat([t.view(-1) for t in vec])) # pylint:disable=not-callable
+        vec_norm = torch.linalg.vector_norm(torch.cat([t.ravel() for t in vec])) # pylint:disable=not-callable
         if vec_norm == 0: return None, [torch.zeros_like(p) for p in params]
         vec = torch._foreach_div(vec, vec_norm)
 

torchzero/utils/linalg/__init__.py

@@ -1,5 +1,12 @@
-from .matrix_funcs import inv_sqrt_2x2, eigvals_func, singular_vals_func, matrix_power_eigh, x_inv
+from . import linear_operator
+from .matrix_funcs import (
+    eigvals_func,
+    inv_sqrt_2x2,
+    matrix_power_eigh,
+    singular_vals_func,
+    x_inv,
+)
 from .orthogonalize import gram_schmidt
 from .qr import qr_householder
+from .solve import cg, nystrom_approximation, nystrom_sketch_and_solve
 from .svd import randomized_svd
-from .solve import cg, nystrom_approximation, nystrom_sketch_and_solve, steihaug_toint_cg