Complete Super Linear

lib/xlayers/super_module.py

@@ -6,11 +6,14 @@ import abc
 import torch.nn as nn
 from enum import Enum
 
+import spaces
+
 
 class SuperRunMode(Enum):
     """This class defines the enumerations for Super Model Running Mode."""
 
     FullModel = "fullmodel"
+    Candidate = "candidate"
     Default = "fullmodel"
 
 
@@ -20,8 +23,23 @@ class SuperModule(abc.ABC, nn.Module):
     def __init__(self):
         super(SuperModule, self).__init__()
         self._super_run_type = SuperRunMode.Default
+        self._abstract_child = None
 
-    @abc.abstractmethod
+    def set_super_run_type(self, super_run_type):
+        def _reset_super_run(m):
+            if isinstance(m, SuperModule):
+                m._super_run_type = super_run_type
+
+        self.apply(_reset_super_run)
+
+    def apply_candiate(self, abstract_child):
+        if not isinstance(abstract_child, spaces.VirtualNode):
+            raise ValueError(
+                "Invalid abstract child program: {:}".format(abstract_child)
+            )
+        self._abstract_child = abstract_child
+
+    @property
     def abstract_search_space(self):
         raise NotImplementedError
 
@@ -29,13 +47,24 @@ class SuperModule(abc.ABC, nn.Module):
     def super_run_type(self):
         return self._super_run_type
 
+    @property
+    def abstract_child(self):
+        return self._abstract_child
+
     @abc.abstractmethod
     def forward_raw(self, *inputs):
+        """Use the largest candidate for forward. Similar to the original PyTorch model."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def forward_candidate(self, *inputs):
         raise NotImplementedError
 
     def forward(self, *inputs):
         if self.super_run_type == SuperRunMode.FullModel:
             return self.forward_raw(*inputs)
+        elif self.super_run_type == SuperRunMode.Candidate:
+            return self.forward_candidate(*inputs)
         else:
             raise ModeError(
                 "Unknown Super Model Run Mode: {:}".format(self.super_run_type)