Merge pull request #38 from LukasHedegaard/develop

Call-mode specific functions in `co.Lambda`

Author: LukasHedegaard

CHANGELOG.md

@@ -9,6 +9,10 @@ From v1.0.0 and on, the project will adherence strictly to Semantic Versioning.
 
 ## [Unreleased]
 
+## [0.15.3]
+### Added
+- Call-mode specific functions in `co.Lambda`
+
 
 ## [0.15.2]
 ### Added

continual/closure.py

@@ -15,25 +15,67 @@ class Lambda(CoModule, nn.Module):
 
     Args:
         fn (Callable[[Tensor], Tensor]): Function to be called during forward.
+        forward_only_fn (Callable[[Tensor], Tensor]): Function to be called only during `forward`. `fn` is used for the other call modes.
+        forward_step_only_fn (Callable[[Tensor], Tensor]): Function to be called only during `forward_step`. `fn` is used for the other call modes.
+        forward_steps_only_fn (Callable[[Tensor], Tensor]): Function to be called only during `forward_steps`. `fn` is used for the other call modes.
+        forward_only_fn (Callable[[Tensor], Tensor]): Function to be called only during forward. `fn` is used for the other call modes.
         takes_time (bool, optional): If True, `fn` recieves all steps, if False, it received one step and no time dimension. Defaults to False.
     """
 
-    def __init__(self, fn: Callable[[Tensor], Tensor], takes_time=False):
+    def __init__(
+        self,
+        fn: Callable[[Tensor], Tensor] = None,
+        forward_only_fn=None,
+        forward_step_only_fn=None,
+        forward_steps_only_fn=None,
+        takes_time=False,
+    ):
         nn.Module.__init__(self)
-        assert callable(fn), "The pased function should be callable."
+        assert callable(fn) or all(
+            callable(forward_only_fn),
+            callable(forward_step_only_fn),
+            callable(forward_steps_only_fn),
+        ), "Either fn or all of forward_only_fn, forward_step_only_fn, and forward_steps_only_fn should be callable."
+
         self.fn = fn
-        if not hasattr(self.fn, "__name__") and hasattr(self.fn, "__repr__"):
-            self.fn.__name__ = self.fn.__repr__()
+        self.forward_only_fn = forward_only_fn
+        self.forward_step_only_fn = forward_step_only_fn
+        self.forward_steps_only_fn = forward_steps_only_fn
         self.takes_time = takes_time
 
+    @staticmethod
+    def build_from(
+        fn: Callable[[Tensor], Tensor],
+        forward_only_fn=None,
+        forward_step_only_fn=None,
+        forward_steps_only_fn=None,
+        takes_time=False,
+    ) -> "Lambda":
+        return Lambda(
+            fn, forward_only_fn, forward_step_only_fn, forward_steps_only_fn, takes_time
+        )
+
     def __repr__(self) -> str:
-        s = f"Lambda({function_repr(self.fn)}"
+        s = "Lambda("
+        if callable(self.fn):
+            s += f"{function_repr(self.fn)}"
+        if callable(self.forward_only_fn):
+            if callable(self.fn):
+                s += ", "
+            s += f"{function_repr(self.forward_only_fn)}"
+        if callable(self.forward_step_only_fn):
+            s += f", {function_repr(self.forward_step_only_fn)}"
+        if callable(self.forward_steps_only_fn):
+            s += f", {function_repr(self.forward_steps_only_fn)}"
         if self.takes_time:
             s += ", takes_time=True"
         s += ")"
         return s
 
     def forward(self, input: Tensor) -> Tensor:
+        if self.forward_only_fn is not None:
+            return self.forward_only_fn(input)
+
         if self.takes_time:
             return self.fn(input)
 
@@ -42,20 +84,27 @@ def forward(self, input: Tensor) -> Tensor:
         )
 
     def forward_steps(self, input: Tensor, pad_end=False, update_state=True) -> Tensor:
-        return self.forward(input)
+        if self.forward_steps_only_fn is not None:
+            return self.forward_steps_only_fn(input)
+
+        if self.takes_time:
+            return self.fn(input)
+
+        return torch.stack(
+            [self.fn(input[:, :, t]) for t in range(input.shape[2])], dim=2
+        )
 
     def forward_step(self, input: Tensor, update_state=True) -> Tensor:
+        if self.forward_step_only_fn is not None:
+            return self.forward_step_only_fn(input)
+
         if self.takes_time:
             input = input.unsqueeze(dim=2)
         output = self.fn(input)
         if self.takes_time:
             output = output.squeeze(dim=2)
         return output
 
-    @staticmethod
-    def build_from(fn: Callable[[Tensor], Tensor], takes_time=False) -> "Lambda":
-        return Lambda(fn, takes_time)
-
 
 def _multiply(x: Tensor, factor: Union[float, int, Tensor]):
     return x * factor

continual/utils.py

@@ -167,6 +167,12 @@ def num_from(tuple_or_num: Union[Number, Tuple[Number, ...]], dim=0) -> Number:
 
 
 def function_repr(fn):
+    if fn is None:
+        return ""
+
+    if isinstance(fn, nn.Module):
+        fn.__name__ = fn.__repr__()
+
     if isinstance(fn, partial):
         fn = fn.func
     s = fn.__name__

setup.py

@@ -25,7 +25,7 @@ def from_file(file_name: str = "requirements.txt", comment_char: str = "#"):
 
 setup(
     name="continual-inference",
-    version="0.15.2",
+    version="0.15.3",
     description="Building blocks for Continual Inference Networks in PyTorch",
     long_description=long_description(),
     long_description_content_type="text/markdown",

tests/continual/test_closure.py

@@ -63,13 +63,12 @@ def local_always42(x):
 
     # Anonymous
     mod = Lambda.build_from(lambda x: torch.ones_like(x) * 42)
-    assert torch.equal(target, mod(x))
+    assert torch.equal(target, mod.forward_steps(x))
 
     # Functor
     functor = torch.nn.Sigmoid()
     assert torch.equal(functor(x), Lambda(functor)(x))
 
-    # takes_time = False
     mod = Lambda.build_from(lambda x: torch.ones_like(x) * 42, takes_time=True)
     assert torch.equal(target, mod(x))
 
@@ -83,6 +82,27 @@ def local_always42(x):
     assert modules[0][1].__repr__() == "Lambda(lambda x: x.view(x.shape[0], -1))"
 
 
+def test_lambda_call_specific_fns():
+    x = torch.ones((1, 1, 2, 2))
+
+    mod = Lambda(
+        fn=lambda x: x + 1,
+        forward_only_fn=lambda x: x + 2,
+        forward_step_only_fn=lambda x: x + 3,
+        forward_steps_only_fn=lambda x: x + 4,
+    )
+
+    assert torch.equal(x + 2, mod.forward(x))
+    assert torch.equal(x.squeeze(2) + 3, mod.forward_step(x))
+    assert torch.equal(x + 4, mod.forward_steps(x))
+
+    # __repr__
+    assert (
+        mod.__repr__()
+        == "Lambda(lambda x: x + 1, lambda x: x + 2, lambda x: x + 3, lambda x: x + 4)"
+    )
+
+
 def test_unity():
     x = torch.ones((1, 1, 2, 2))
     assert torch.equal(x, Unity()(x))