Add output shape to pooling layers

Author: APJansen

groco/layers/pooling.py

@@ -76,6 +76,13 @@ def get_config(self):
         config.update(self.group_transforms.get_config())
         return config
 
+    def compute_output_shape(self, input_shape):
+        ga = self.group_transforms.group_axis
+        input_shape_normal = input_shape[:ga] + input_shape[ga + 1 :]
+        output_shape = self.pooling.compute_output_shape(input_shape_normal)
+        go = self.group.order
+        return output_shape[:ga] + (go,) + output_shape[ga:]
+
 
 class GlobalGroupPooling(Layer):
     """
@@ -127,6 +134,12 @@ def get_config(self):
         config = self.pooling.get_config()
         return config
 
+    def compute_output_shape(self, input_shape):
+        ga = self.group_transforms.group_axis
+        input_shape_normal = input_shape[:ga] + input_shape[ga + 1 :]
+        output_shape = self.pooling.compute_output_shape(input_shape_normal)
+        return output_shape[:ga] + output_shape[ga:]
+
 
 class GroupMaxPooling1D(GroupPooling):
     """