Add pack layer tests, fix unpack_3d layer

packnet_tests.py

@@ -0,0 +1,33 @@
+import unittest
+
+import tensorflow as tf
+
+import packnet_functional as p
+
+
+class PacknetTests(unittest.TestCase):
+
+    def test_pack_3d_layer(self):
+        # 3d packing expects a multiple of 16 for channels due to using 16 groups in group normalisation
+        test_input = tf.random.normal([4, 224, 224, 32])
+        y = p.pack_3d(test_input, 3, features_3d=4)
+        out_shape = [i for i in test_input.shape]
+        out_shape[1] = out_shape[1] // 2
+        out_shape[2] = out_shape[2] // 2
+        # TODO: Anything else we can test here for validity?
+        self.assertEqual(y.shape, out_shape)
+
+    def test_unpack_3d_layer(self):
+        num_output_channels = 32
+        test_input = tf.random.normal([4, 112, 112, 64])
+        y = p.unpack_3d(test_input, num_output_channels, 3, features_3d=4)
+        out_shape = [i for i in test_input.shape]
+        out_shape[1] = out_shape[1] * 2
+        out_shape[2] = out_shape[2] * 2
+        out_shape[3] = num_output_channels
+        # TODO: Anything else we can test here for validity?
+        self.assertEqual(y.shape, out_shape)
+
+
+if __name__ == '__main__':
+    unittest.main()