Add StatefulTupleSelect layer

Signed-off-by: Emanuele Ballarin <[email protected]>

ebtorch/__init__.py

@@ -93,6 +93,7 @@
 del SwiGLU
 del TupleDecouple
 del TupleSelect
+del StatefulTupleSelect
 del SilhouetteScore
 del Concatenate
 del DuplexLinearNeck

ebtorch/nn/__init__.py

@@ -38,6 +38,7 @@
 from .architectures import SharedDuplexLinearNeck
 from .architectures import SilhouetteScore
 from .architectures import SirenSine
+from .architectures import StatefulTupleSelect
 from .architectures import SwiGLU
 from .architectures import TupleDecouple
 from .architectures import TupleSelect

ebtorch/nn/architectures.py

@@ -61,6 +61,7 @@
     "SwiGLU",
     "TupleDecouple",
     "TupleSelect",
+    "StatefulTupleSelect",
     "SilhouetteScore",
     "Concatenate",
     "DuplexLinearNeck",
@@ -917,6 +918,16 @@ def forward(self, xtuple: Tuple[Tensor, ...]) -> Tensor:
         return xtuple[self.idx]
 
 
+class StatefulTupleSelect(nn.Module):
+    def __init__(self, module, idx: int = 0) -> None:
+        super().__init__()
+        self.module: nn.Module = module
+        self.idx: int = idx
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.module(x)[self.idx]
+
+
 class SilhouetteScore(nn.Module):
     """
     Layerized computation of the Silhouette Score.

setup.py

@@ -26,7 +26,7 @@ def read(fname):
 
 setup(
     name=PACKAGENAME,
-    version="0.25.12",
+    version="0.25.13",
     author="Emanuele Ballarin",
     author_email="[email protected]",
     url="https://github.com/emaballarin/ebtorch",