Add LambdaLayer

prototorch/modules/__init__.py

@@ -1 +1,3 @@
-"""ProtoTorch modules."""
+"""ProtoTorch modules."""
+
+from .utils import LambdaLayer

prototorch/modules/utils.py

@@ -0,0 +1,16 @@
+"""ProtoTorch utility modules."""
+
+import torch
+
+
+class LambdaLayer(torch.nn.Module):
+    def __init__(self, fn, name=None):
+        super().__init__()
+        self.fn = fn
+        self.name = name or fn.__name__ # lambda fns get <lambda>
+
+    def forward(self, *args, **kwargs):
+        return self.fn(*args, **kwargs)
+
+    def extra_repr(self):
+        return self.name