Update super-activation layers

lib/xlayers/super_module.py

@@ -11,128 +11,10 @@ from enum import Enum
 
 import spaces
 
-IntSpaceType = Union[int, spaces.Integer, spaces.Categorical]
-BoolSpaceType = Union[bool, spaces.Categorical]
-
-
-class LayerOrder(Enum):
-    """This class defines the enumerations for order of operation in a residual or normalization-based layer."""
-
-    PreNorm = "pre-norm"
-    PostNorm = "post-norm"
-
-
-class SuperRunMode(Enum):
-    """This class defines the enumerations for Super Model Running Mode."""
-
-    FullModel = "fullmodel"
-    Candidate = "candidate"
-    Default = "fullmodel"
-
-
-class TensorContainer:
-    """A class to maintain both parameters and buffers for a model."""
-
-    def __init__(self):
-        self._names = []
-        self._tensors = []
-        self._param_or_buffers = []
-        self._name2index = dict()
-
-    def additive(self, tensors):
-        result = TensorContainer()
-        for index, name in enumerate(self._names):
-            new_tensor = self._tensors[index] + tensors[index]
-            result.append(name, new_tensor, self._param_or_buffers[index])
-        return result
-
-    def create_container(self, tensors):
-        result = TensorContainer()
-        for index, name in enumerate(self._names):
-            new_tensor = tensors[index]
-            result.append(name, new_tensor, self._param_or_buffers[index])
-        return result
-
-    def no_grad_clone(self):
-        result = TensorContainer()
-        with torch.no_grad():
-            for index, name in enumerate(self._names):
-                result.append(
-                    name, self._tensors[index].clone(), self._param_or_buffers[index]
-                )
-        return result
-
-    def requires_grad_(self, requires_grad=True):
-        for tensor in self._tensors:
-            tensor.requires_grad_(requires_grad)
-
-    def parameters(self):
-        return self._tensors
-
-    @property
-    def tensors(self):
-        return self._tensors
-
-    def flatten(self, tensors=None):
-        if tensors is None:
-            tensors = self._tensors
-        tensors = [tensor.view(-1) for tensor in tensors]
-        return torch.cat(tensors)
-
-    def unflatten(self, tensor):
-        tensors, s = [], 0
-        for raw_tensor in self._tensors:
-            length = raw_tensor.numel()
-            x = torch.reshape(tensor[s : s + length], shape=raw_tensor.shape)
-            tensors.append(x)
-            s += length
-        return tensors
-
-    def append(self, name, tensor, param_or_buffer):
-        if not isinstance(tensor, torch.Tensor):
-            raise TypeError(
-                "The input tensor must be torch.Tensor instead of {:}".format(
-                    type(tensor)
-                )
-            )
-        self._names.append(name)
-        self._tensors.append(tensor)
-        self._param_or_buffers.append(param_or_buffer)
-        assert name not in self._name2index, "The [{:}] has already been added.".format(
-            name
-        )
-        self._name2index[name] = len(self._names) - 1
-
-    def query(self, name):
-        if not self.has(name):
-            raise ValueError(
-                "The {:} is not in {:}".format(name, list(self._name2index.keys()))
-            )
-        index = self._name2index[name]
-        return self._tensors[index]
-
-    def has(self, name):
-        return name in self._name2index
-
-    def has_prefix(self, prefix):
-        for name, idx in self._name2index.items():
-            if name.startswith(prefix):
-                return name
-        return False
-
-    def numel(self):
-        total = 0
-        for tensor in self._tensors:
-            total += tensor.numel()
-        return total
-
-    def __len__(self):
-        return len(self._names)
-
-    def __repr__(self):
-        return "{name}({num} tensors)".format(
-            name=self.__class__.__name__, num=len(self)
-        )
+from .super_utils import IntSpaceType, BoolSpaceType
+from .super_utils import LayerOrder, SuperRunMode
+from .super_utils import TensorContainer
+from .super_utils import ShapeContainer
 
 
 class SuperModule(abc.ABC, nn.Module):