use namedtuple for embedding wrapper return, for clarity

Author: lucidrains

perfusion_pytorch/embedding.py

@@ -2,11 +2,20 @@
 from torch import nn
 from torch.nn import Module
 
+from collections import namedtuple
+
 from beartype import beartype
 from beartype.typing import Optional, Tuple, Union
 
 from einops import rearrange
 
+# constants
+
+EmbeddingReturn = namedtuple('EmbeddingReturn', [
+    'embed_with_concept',
+    'embed_with_superclass'
+])
+
 # helper functions
 
 def exists(val):
@@ -63,7 +72,7 @@ def forward(
         x,
         concept_id: Optional[Union[int, Tuple[int, ...]]] = None,
         return_embed_with_superclass = True
-    ):
+    ) -> EmbeddingReturn:
         concept_masks = tuple(concept_id == x for concept_id in self.concept_embed_ids)
 
         if exists(concept_id):
@@ -101,9 +110,9 @@ def forward(
             with torch.no_grad():
                 superclass_embeds = self.embed(x)
 
-            return embeds, superclass_embeds
+            return EmbeddingReturn(embeds, superclass_embeds)
 
-        return embeds
+        return EmbeddingReturn(embeds, None)
 
 @beartype
 def merge_embedding_wrappers(

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'perfusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.8',
+  version = '0.1.9',
   license='MIT',
   description = 'Perfusion - Pytorch',
   author = 'Phil Wang',