refactor prepare_call

Author: APJansen

groco/layers/conv2d.py

@@ -60,12 +60,11 @@ def __init__(
         self.subgroup = self.group_transforms.subgroup
 
         self.group_valued_input = None
-        self.group_order = self.group.order if transpose else self.subgroup.order
 
     @backup_and_restore(("kernel", "bias", "filters"))
     def call(self, inputs):
         self.kernel, self.bias, self.filters, inputs = self.group_transforms.prepare_call(
-            self.kernel, self.bias, self.filters, inputs, self.use_bias, self.group_order
+            self.kernel, self.bias, self.filters, inputs, self.use_bias
         )
         outputs = super().call(inputs)
         return self.group_transforms.restore_group_axis(outputs)

groco/layers/conv2d_transpose.py

@@ -62,7 +62,7 @@ def __init__(
     @backup_and_restore(("kernel", "bias", "filters"))
     def call(self, inputs):
         self.kernel, self.bias, self.filters, inputs = self.group_transforms.prepare_call(
-            self.kernel, self.bias, self.filters, inputs, self.use_bias, self.group_order
+            self.kernel, self.bias, self.filters, inputs, self.use_bias
         )
         outputs = super().call(inputs)
         return self.group_transforms.restore_group_axis(outputs)

groco/layers/conv3d.py

@@ -65,7 +65,7 @@ def __init__(
     @backup_and_restore(("kernel", "bias", "filters"))
     def call(self, inputs):
         self.kernel, self.bias, self.filters, inputs = self.group_transforms.prepare_call(
-            self.kernel, self.bias, self.filters, inputs, self.use_bias, self.group_order
+            self.kernel, self.bias, self.filters, inputs, self.use_bias
         )
         outputs = super().call(inputs)
         return self.group_transforms.restore_group_axis(outputs)

groco/layers/conv3d_transpose.py

@@ -68,7 +68,7 @@ def __init__(
     @backup_and_restore(("kernel", "bias", "filters"))
     def call(self, inputs):
         self.kernel, self.bias, self.filters, inputs = self.group_transforms.prepare_call(
-            self.kernel, self.bias, self.filters, inputs, self.use_bias, self.group_order
+            self.kernel, self.bias, self.filters, inputs, self.use_bias
         )
         outputs = super().call(inputs)
         return self.group_transforms.restore_group_axis(outputs)

groco/layers/group_transforms.py

@@ -243,10 +243,11 @@ def get_config(self):
         config.update(self.equivariant_padding.get_config())
         return config
 
-    def prepare_call(self, kernel, bias, filters, inputs, use_bias, group_order):
+    def prepare_call(self, kernel, bias, filters, inputs, use_bias):
         inputs = self.merge_group_axis_and_pad(inputs)
         kernel = self.transform_kernel(kernel)
         if use_bias:
             bias = self.repeat_bias(bias)
-        filters *= group_order
+        factor = len(self.group.subgroup[self.acting_group])
+        filters *= factor
         return kernel, bias, filters, inputs