Abstracted away reordering.

Author: vbharadwaj-bk

openequivariance/openequivariance/core/ComputationSchedule.py

@@ -619,40 +619,33 @@ def calculate_backward_smem(
             smem=self.memory_per_warp * warps_per_block,
         )
 
-    def reorder_weights(self, weights_in, direction, has_batch_dim):
+    def weight_reordering_info(self, weights_in, has_batch_dim):
         """
-        Reorders weights from the canonical e3nn form to the
-        form that LoopUnrollTP can ingest. Can also reorder the parameters
-        of a dense neural network layer that produces the weight matrix.
-
-        If has_batch_dim is true, the first dimension of the input weight matrix
-        is treated as the batch dimension.
+        Calculates all shapes, slices, and permutation info to reorder
+        weights. 
         """
-        import torch  # TODO-someday: no need to specialize this to PyTorch
+        batch_dim = weights_in.shape[0]
+        reorder_specs = []
 
-        weights_out = torch.zeros_like(weights_in)
-        assert direction in ["forward", "backward"]
         for i, child_inst in enumerate(self.problem_splitter.new_instructions):
             parent_start, parent_end = (
                 child_inst.parent_weights_start,
                 child_inst.parent_weights_end,
             )
             parent_shape = list(child_inst.parent_weights_shape)
+            parent_range = [slice(parent_start, parent_end)]
 
             child_start, child_end, child_shape = (
                 self.updated_config.weight_range_and_shape_for_instruction(i)
             )
-
-            parent_range, child_range = (
-                [slice(parent_start, parent_end)],
-                [slice(child_start, child_end)],
-            )
+            child_range = [slice(child_start, child_end)]
+            
             weights_subrange = child_inst.weights_subrange
-            batch_dim = weights_in.shape[0]
+            
             reshape_size = [-1]
             transpose_perm = None
-
             connection_mode = self.updated_config.instructions[i].connection_mode
+            
             if connection_mode == "uvu":
                 transpose_perm = [1, 0]
             elif connection_mode == "uvw":
@@ -662,50 +655,27 @@ def reorder_weights(self, weights_in, direction, has_batch_dim):
                 child_range = [slice(0, batch_dim)] + child_range
                 parent_range = [slice(0, batch_dim)] + parent_range
                 parent_shape = [batch_dim] + parent_shape
+                
                 child_shape = [batch_dim] + list(child_shape)
                 weights_subrange = [slice(0, batch_dim)] + child_inst.weights_subrange
                 reshape_size = [batch_dim] + reshape_size
-                transpose_perm = [0] + [i + 1 for i in transpose_perm]
-
-            if direction == "forward":
-                sliced_weights = weights_in[tuple(parent_range)].reshape(parent_shape)[
-                    tuple(weights_subrange)
-                ]
-                weights_out[tuple(child_range)] = sliced_weights.permute(
-                    transpose_perm
-                ).reshape(reshape_size)
-            elif direction == "backward":
-                transpose_child_shape = [child_shape[i] for i in transpose_perm]
-                sliced_weights = (
-                    weights_in[tuple(child_range)]
-                    .reshape(transpose_child_shape)
-                    .permute(transpose_perm)
-                )
-                weights_out[tuple(parent_range)].reshape(parent_shape)[
-                    tuple(weights_subrange)
-                ] = sliced_weights.flatten().reshape(child_shape)
-
-        return weights_out
-
-    def reorder_weights_numpy(self, weights_in, direction, has_batch_dim):
-        import torch
-
-        weights_in = torch.from_numpy(weights_in.copy())
-        result = self.reorder_weights(weights_in, direction, has_batch_dim)
-        return result.detach().cpu().numpy().copy()
-
-    def reorder_weights_from_e3nn(self, weights_in, has_batch_dim):
-        import torch
-
-        if isinstance(weights_in, np.ndarray):
-            return self.reorder_weights_numpy(weights_in, "forward", has_batch_dim)
-        elif isinstance(weights_in, torch.Tensor):
-            return self.reorder_weights(weights_in, "forward", has_batch_dim)
-
-    def reorder_weights_to_e3nn(self, weights_in, has_batch_dim):
-        import torch
-
-        if isinstance(weights_in, np.ndarray):
-            return self.reorder_weights_numpy(weights_in, "backward", has_batch_dim)
-        elif isinstance(weights_in, torch.Tensor):
-            return self.reorder_weights(weights_in, "backward", has_batch_dim)
+                
+                if transpose_perm is not None:
+                    transpose_perm = [0] + [k + 1 for k in transpose_perm]
+
+            transpose_child_shape = None
+            if transpose_perm is not None:
+                transpose_child_shape = [child_shape[k] for k in transpose_perm]
+
+            reorder_specs.append({
+                "parent_range": tuple(parent_range),
+                "parent_shape": parent_shape,
+                "weights_subrange": tuple(weights_subrange),
+                "child_range": tuple(child_range),
+                "child_shape": child_shape,
+                "transpose_perm": transpose_perm,
+                "reshape_size": reshape_size,
+                "transpose_child_shape": transpose_child_shape,
+            })
+
+        return reorder_specs

openequivariance/openequivariance/impl_torch/TensorProduct.py

@@ -5,6 +5,8 @@
 import typing
 from openequivariance.core.utils import torch_to_oeq_dtype
 from openequivariance.benchmark.logging_utils import getLogger
+from openequivariance.impl_torch.utils import reorder_torch
+
 
 logger = getLogger()
 
@@ -90,10 +92,10 @@ def __setstate__(self, state):
         self._init_class()
 
     def reorder_weights_from_e3nn(self, weights, has_batch_dim=True):
-        return self.forward_schedule.reorder_weights_from_e3nn(weights, has_batch_dim)
+        return reorder_torch(self.forward_schedule, weights, "forward", not self.config.shared_weights)
 
     def reorder_weights_to_e3nn(self, weights, has_batch_dim=True):
-        return self.forward_schedule.reorder_weights_to_e3nn(weights, has_batch_dim)
+        return reorder_torch(self.forward_schedule, weights, "backward", not self.config.shared_weights)
 
     def forward(
         self, x: torch.Tensor, y: torch.Tensor, W: torch.Tensor

openequivariance/openequivariance/impl_torch/TensorProductConv.py

@@ -19,6 +19,7 @@
 from openequivariance import TPProblem
 from openequivariance.core.utils import torch_to_oeq_dtype
 from openequivariance.core.dtype_enum import enum_to_torch_dtype
+from openequivariance.impl_torch.utils import reorder_torch
 
 from openequivariance.benchmark.logging_utils import getLogger
 
@@ -418,10 +419,10 @@ def double_backward(ctx, grad_output):
         )
 
     def reorder_weights_from_e3nn(self, weights, has_batch_dim=True):
-        return self.forward_schedule.reorder_weights_from_e3nn(weights, has_batch_dim)
+        return reorder_torch(self.forward_schedule, weights, "forward", not self.config.shared_weights)
 
     def reorder_weights_to_e3nn(self, weights, has_batch_dim=True):
-        return self.forward_schedule.reorder_weights_to_e3nn(weights, has_batch_dim)
+        return reorder_torch(self.forward_schedule, weights, "backward", not self.config.shared_weights)
 
     @staticmethod
     def name():

openequivariance/openequivariance/impl_torch/utils.py

@@ -0,0 +1,52 @@
+import torch
+
+def reorder_helper(schedule, weights_in, direction, has_batch_dim):
+        assert direction in ["forward", "backward"]
+
+        specs = schedule.weight_reordering_info(weights_in, has_batch_dim) 
+        weights_out = torch.zeros_like(weights_in)
+
+        for spec in specs:
+            parent_range = spec["parent_range"]
+            parent_shape = spec["parent_shape"]
+            weights_subrange = spec["weights_subrange"]
+            child_range = spec["child_range"]
+            transpose_perm = spec["transpose_perm"]
+            
+            if direction == "forward":
+                reshape_size = spec["reshape_size"]
+                
+                sliced_weights = weights_in[parent_range].reshape(parent_shape)[
+                    weights_subrange
+                ]
+                
+                weights_out[child_range] = sliced_weights.permute(
+                    transpose_perm
+                ).reshape(reshape_size)
+
+            elif direction == "backward":
+                transpose_child_shape = spec["transpose_child_shape"]
+                child_shape = spec["child_shape"]
+
+                sliced_weights = (
+                    weights_in[child_range]
+                    .reshape(transpose_child_shape)
+                    .permute(transpose_perm)
+                )
+
+                weights_out[parent_range].reshape(parent_shape)[
+                    weights_subrange
+                ] = sliced_weights.flatten().reshape(child_shape)
+
+        return weights_out
+
+def reorder_numpy_helper(schedule, weights_in, direction, has_batch_dim):
+    weights_in = torch.from_numpy(weights_in.copy())
+    result = reorder_helper(schedule, weights_in, direction, has_batch_dim)
+    return result.detach().cpu().numpy().copy()
+
+def reorder_torch(schedule, weights_in, direction, has_batch_dim):
+    if isinstance(weights_in, torch.Tensor):
+        return reorder_helper(schedule, weights_in, direction, has_batch_dim)
+    else:
+        return reorder_numpy_helper(schedule, weights_in, direction, has_batch_dim)