Complete xlayers.rearrange

tests/test_super_att.py

@@ -29,6 +29,7 @@ class TestSuperSelfAttention(unittest.TestCase):
         abstract_child = abstract_space.random(reuse_last=True)
         print("The abstract child program is:\n{:}".format(abstract_child))
         model.set_super_run_type(super_core.SuperRunMode.Candidate)
+        model.enable_candidate()
         model.apply_candidate(abstract_child)
         outputs = model(inputs)
         return abstract_child, outputs

tests/test_super_container.py

@@ -25,6 +25,7 @@ def _internal_func(inputs, model):
     abstract_space.clean_last()
     abstract_child = abstract_space.random(reuse_last=True)
     print("The abstract child program is:\n{:}".format(abstract_child))
+    model.enable_candidate()
     model.set_super_run_type(super_core.SuperRunMode.Candidate)
     model.apply_candidate(abstract_child)
     outputs = model(inputs)

tests/test_super_mlp.py

@@ -37,6 +37,7 @@ class TestSuperLinear(unittest.TestCase):
         print("The abstract child program:\n{:}".format(abstract_child))
 
         model.set_super_run_type(super_core.SuperRunMode.Candidate)
+        model.enable_candidate()
         model.apply_candidate(abstract_child)
 
         output_shape = (20, abstract_child["_out_features"].value)
@@ -77,6 +78,7 @@ class TestSuperLinear(unittest.TestCase):
         )
 
         mlp.set_super_run_type(super_core.SuperRunMode.Candidate)
+        mlp.enable_candidate()
         mlp.apply_candidate(abstract_child)
         outputs = mlp(inputs)
         output_shape = (4, abstract_child["fc2"]["_out_features"].value)
@@ -103,6 +105,7 @@ class TestSuperLinear(unittest.TestCase):
         print("The abstract child program is:\n{:}".format(abstract_child))
 
         mlp.set_super_run_type(super_core.SuperRunMode.Candidate)
+        mlp.enable_candidate()
         mlp.apply_candidate(abstract_child)
         outputs = mlp(inputs)
         output_shape = (4, abstract_child["_out_features"].value)
@@ -120,6 +123,7 @@ class TestSuperLinear(unittest.TestCase):
         print("The abstract child program:\n{:}".format(abstract_child))
 
         model.set_super_run_type(super_core.SuperRunMode.Candidate)
+        model.enable_candidate()
         model.apply_candidate(abstract_child)
         outputs = model(inputs)
         output_shape = (4, 60, abstract_child["_embed_dim"].value)

tests/test_super_norm.py

@@ -38,6 +38,7 @@ class TestSuperSimpleNorm(unittest.TestCase):
         print("The abstract child program:\n{:}".format(abstract_child))
 
         model.set_super_run_type(super_core.SuperRunMode.Candidate)
+        model.enable_candidate()
         model.apply_candidate(abstract_child)
 
         output_shape = (20, abstract_child["1"]["_out_features"].value)
@@ -70,6 +71,7 @@ class TestSuperSimpleNorm(unittest.TestCase):
         print("The abstract child program:\n{:}".format(abstract_child))
 
         model.set_super_run_type(super_core.SuperRunMode.Candidate)
+        model.enable_candidate()
         model.apply_candidate(abstract_child)
 
         output_shape = (20, abstract_child["2"]["_out_features"].value)

tests/test_super_rearrange.py

@@ -5,12 +5,6 @@
 #####################################################
 import sys
 import unittest
-from pathlib import Path
-
-lib_dir = (Path(__file__).parent / "..").resolve()
-print("LIB-DIR: {:}".format(lib_dir))
-if str(lib_dir) not in sys.path:
-    sys.path.insert(0, str(lib_dir))
 
 import torch
 from xautodl import xlayers
@@ -28,3 +22,4 @@ class TestSuperReArrange(unittest.TestCase):
         print(layer)
         outs = layer(tensor)
         print("The output tensor shape: {:}".format(outs.shape))
+        assert tuple(outs.shape) == (8, 32 * 32 // 16, 4 * 4 * 4)

tests/test_super_transformer.py

@@ -1,36 +0,0 @@
-#####################################################
-# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
-#####################################################
-# pytest ./tests/test_super_model.py -s             #
-#####################################################
-import unittest
-
-import torch
-from xautodl.xlayers.super_core import SuperRunMode
-from xautodl.trade_models import get_transformer
-
-
-class TestSuperTransformer(unittest.TestCase):
-    """Test the super transformer."""
-
-    def test_super_transformer(self):
-        model = get_transformer(None)
-        model.apply_verbose(False)
-        print(model)
-
-        inputs = torch.rand(10, 360)
-        print("Input shape: {:}".format(inputs.shape))
-        outputs = model(inputs)
-        self.assertEqual(tuple(outputs.shape), (10,))
-
-        abstract_space = model.abstract_search_space
-        abstract_space.clean_last()
-        abstract_child = abstract_space.random(reuse_last=True)
-        print("The abstract searc space:\n{:}".format(abstract_space))
-        print("The abstract child program:\n{:}".format(abstract_child))
-
-        model.set_super_run_type(SuperRunMode.Candidate)
-        model.apply_candidate(abstract_child)
-
-        outputs = model(inputs)
-        self.assertEqual(tuple(outputs.shape), (10,))