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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-einops-utils
-Version: 0.1.2
+Version: 0.1.4
 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.4}/pyproject.toml

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

{torch_einops_utils-0.1.2 → torch_einops_utils-0.1.4}/tests/test_nn.py

@@ -1,6 +1,6 @@
 import torch
 from torch import nn
-from torch_einops_utils.nn import Sequential, Lambda, Identity
+from torch_einops_utils.nn import Sequential, Lambda, Identity, Residual
 
 def test_sequential():
     # Test that it filters out None
@@ -26,3 +26,18 @@ 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.]))

{torch_einops_utils-0.1.2 → torch_einops_utils-0.1.4}/torch_einops_utils/nn.py

@@ -2,7 +2,7 @@ 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
+from torch_einops_utils.torch_einops_utils import exists, tree_flatten_with_inverse
 
 def Sequential(*modules):
     return PyTorchSequential(*filter(exists, modules))
@@ -18,3 +18,14 @@ class Lambda(Module):
 
     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))