torch-einops-utils 0.1.2__tar.gz → 0.1.5__tar.gz

{torch_einops_utils-0.1.2 → torch_einops_utils-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-einops-utils
-Version: 0.1.2
+Version: 0.1.5
 Summary: Personal utility functions
 Project-URL: Homepage, https://pypi.org/project/torch-einops-utils/
 Project-URL: Repository, https://github.com/lucidrains/torch-einops-utils

{torch_einops_utils-0.1.2 → torch_einops_utils-0.1.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "torch-einops-utils"
-version = "0.1.2"
+version = "0.1.5"
 description = "Personal utility functions"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

torch_einops_utils-0.1.5/tests/test_nn.py

@@ -0,0 +1,75 @@
+import torch
+from torch import nn
+from torch_einops_utils.nn import Sequential, Lambda, Identity, Residual, count_parameters
+
+def test_sequential():
+    # Test that it filters out None
+    seq = Sequential(
+        nn.Linear(10, 10),
+        None,
+        nn.ReLU()
+    )
+    assert len(seq) == 2
+
+    # Test forward pass
+    x = torch.randn(2, 10)
+    out = seq(x)
+    assert out.shape == (2, 10)
+
+def test_lambda():
+    fn = lambda x: x * 2
+    lam = Lambda(fn)
+    x = torch.tensor([1., 2., 3.])
+    assert torch.allclose(lam(x), torch.tensor([2., 4., 6.]))
+
+def test_identity():
+    ident = Identity()
+    x = torch.tensor([1., 2., 3.])
+    assert torch.allclose(ident(x), x)
+
+def test_residual():
+    fn = lambda x: x * 2
+    res = Residual(fn)
+    x = torch.tensor([1., 2., 3.])
+    assert torch.allclose(res(x), torch.tensor([3., 6., 9.]))
+
+    def fn_tuple(x):
+        return x * 2, x * 3, dict(a = x * 4)
+
+    res_tuple = Residual(fn_tuple)
+    out1, out2, out3 = res_tuple(x)
+    assert torch.allclose(out1, torch.tensor([3., 6., 9.]))
+    assert torch.allclose(out2, torch.tensor([3., 6., 9.]))
+    assert torch.allclose(out3['a'], torch.tensor([4., 8., 12.]))
+
+def test_count_parameters():
+    model = nn.Linear(10, 10)
+    assert count_parameters(model) == 110
+
+    # Test requires_grad filter
+    model.bias.requires_grad_(False)
+    assert count_parameters(model) == 110
+    assert count_parameters(model, requires_grad=True) == 100
+    assert count_parameters(model, requires_grad=False) == 10
+
+    # Test as a decorator
+    @count_parameters
+    class MyModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = nn.Linear(10, 10)
+            self.linear.bias.requires_grad_(False)
+
+    my_model = MyModel()
+    assert my_model.num_parameters == 110
+
+    # Test as a decorator with kwargs
+    @count_parameters(requires_grad=True)
+    class MyModelTrainable(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = nn.Linear(10, 10)
+            self.linear.bias.requires_grad_(False)
+
+    my_model_trainable = MyModelTrainable()
+    assert my_model_trainable.num_parameters == 100

torch_einops_utils-0.1.5/torch_einops_utils/nn.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from typing import Callable
+from functools import partial
+
+from torch.nn import Module
+from torch.nn import Sequential as PyTorchSequential
+
+from torch_einops_utils.torch_einops_utils import exists, tree_flatten_with_inverse
+
+# helpers / functions
+
+def count_parameters(model_or_class = None, *, requires_grad: bool | None = None):
+    if not exists(model_or_class):
+        return partial(count_parameters, requires_grad = requires_grad)
+
+    def _count(model):
+        return sum(p.numel() for p in model.parameters() if not exists(requires_grad) or p.requires_grad == requires_grad)
+
+    if isinstance(model_or_class, type) and issubclass(model_or_class, Module):
+        model_or_class.num_parameters = property(_count)
+        return model_or_class
+
+    return _count(model_or_class)
+
+def Sequential(*modules):
+    return PyTorchSequential(*filter(exists, modules))
+
+# classes
+
+class Identity(Module):
+    def forward(self, t, *args, **kwargs):
+        return t
+
+class Lambda(Module):
+    def __init__(self, fn: Callable):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, t, *args, **kwargs):
+        return self.fn(t, *args, **kwargs)
+
+class Residual(Module):
+    def __init__(self, fn: Callable):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x, *args, **kwargs):
+        out = self.fn(x, *args, **kwargs)
+
+        (first, *rest), inverse = tree_flatten_with_inverse(out)
+        return inverse((first + x, *rest))

{torch_einops_utils-0.1.2 → torch_einops_utils-0.1.5}/torch_einops_utils/save_load.py

@@ -67,8 +67,10 @@ def save_load(
             setattr(self, config_instance_var_name, (args, kwargs))
             _orig_init(self, *args, **kwargs)
 
-        def _save(self, path, overwrite = True):
-            path = Path(path)
+        def _save(self, path: str | Path, overwrite = True):
+            if isinstance(path, str):
+                path = Path(path)
+
             assert overwrite or not path.exists()
 
             config = getattr(self, config_instance_var_name)
@@ -80,8 +82,10 @@ def save_load(
 
             torch.save(pkg, str(path))
 
-        def _load(self, path, strict = True):
-            path = Path(path)
+        def _load(self, path: str | Path, strict = True):
+            if isinstance(path, str):
+                path = Path(path)
+
             assert path.exists()
 
             pkg = torch.load(str(path), map_location = 'cpu')
@@ -95,8 +99,10 @@ def save_load(
         # looks for a `config` key in the stored checkpoint, instantiating the model as well as loading the state dict
 
         @classmethod
-        def _init_and_load_from(cls, path, strict = True):
-            path = Path(path)
+        def _init_and_load_from(cls, path: str | Path, strict = True):
+            if isinstance(path, str):
+                path = Path(path)
+
             assert path.exists()
             pkg = torch.load(str(path), map_location = 'cpu')
 

torch_einops_utils-0.1.2/tests/test_nn.py

@@ -1,28 +0,0 @@
-import torch
-from torch import nn
-from torch_einops_utils.nn import Sequential, Lambda, Identity
-
-def test_sequential():
-    # Test that it filters out None
-    seq = Sequential(
-        nn.Linear(10, 10),
-        None,
-        nn.ReLU()
-    )
-    assert len(seq) == 2
-
-    # Test forward pass
-    x = torch.randn(2, 10)
-    out = seq(x)
-    assert out.shape == (2, 10)
-
-def test_lambda():
-    fn = lambda x: x * 2
-    lam = Lambda(fn)
-    x = torch.tensor([1., 2., 3.])
-    assert torch.allclose(lam(x), torch.tensor([2., 4., 6.]))
-
-def test_identity():
-    ident = Identity()
-    x = torch.tensor([1., 2., 3.])
-    assert torch.allclose(ident(x), x)

torch_einops_utils-0.1.2/torch_einops_utils/nn.py

@@ -1,20 +0,0 @@
-from __future__ import annotations
-from typing import Callable
-
-from torch.nn import Module, Sequential as PyTorchSequential
-from torch_einops_utils.torch_einops_utils import exists
-
-def Sequential(*modules):
-    return PyTorchSequential(*filter(exists, modules))
-
-class Identity(Module):
-    def forward(self, t, *args, **kwargs):
-        return t
-
-class Lambda(Module):
-    def __init__(self, fn: Callable):
-        super().__init__()
-        self.fn = fn
-
-    def forward(self, t, *args, **kwargs):
-        return self.fn(t, *args, **kwargs)