Feat (equalize): adding initial support for MixQuant

src/brevitas/graph/equalize.py

@@ -39,7 +39,6 @@
 from brevitas.graph.base import Transform
 from brevitas.graph.hadamard import find_closest_hadamard_number
 from brevitas.graph.hadamard import get_hadK
-from brevitas.graph.hadamard import is_pow2
 from brevitas.graph.hadamard import matmul_hadU
 from brevitas.graph.hadamard import matmul_hadU_cuda
 from brevitas.graph.hadamard import random_hadamard_matrix
@@ -720,6 +719,18 @@ def from_module_indexes(
 
         return cls(module, weight_axis, act_axis, indexes)
 
+    def permute(self, permute_index):
+        self.module.weight.data = torch.index_select(
+            self.module.weight.data, self.weight_axis, permute_index.to(self.module.weight.device))
+        if hasattr(self.module, self._bias_tensor_name):
+            bias = getattr(self.module, self._bias_tensor_name)
+            # hasattr returns true if bias=None
+            if bias is not None:
+                bias.data = torch.index_select(
+                    self.module.bias.data,
+                    self.weight_axis,
+                    permute_index.to(self.module.bias.device))
+
 
 class EqualizationSinkWrapper(EqualizationModuleWrapper):
 
@@ -760,6 +771,10 @@ def from_module_indexes(
             weight_tensor_name = "weight"
         return cls(module, weight_axis, act_axis, indexes, weight_tensor_name)
 
+    def permute(self, permute_index):
+        self.module.weight.data = torch.index_select(
+            self.module.weight.data, self.weight_axis, permute_index.to(self.module.weight.device))
+
 
 # When fuse_scaling = False, the scaling parameters are instances of nn.Parameter,
 # which are registered to the scaling modules (used in the parametrization of the
@@ -2060,6 +2075,8 @@ def find_sink(node):
             end_index = head_dim if head_dim != -1 else output_weight.shape[0]
             output_index = EqualizationIndexes(0, end_index, 0)
 
+            # NOTE: GraphPermutationEqualization.extract_permute_regions looks for these src and
+            # sink names to delineate SDPA regions
             region = Region.from_dicts(
                 srcs={'value_sdpa': value_index},
                 sinks={'output_sdpa': output_index},
@@ -2128,6 +2145,9 @@ def apply(self,
                 added_regions += 1
         logging.debug(f"Adding {added_regions} sink-only regions")
 
+        # Store regions for potential use by GraphPermutationEqualization
+        self.regions = regions
+
         if overlap:
             assert not self.use_parametrized_rotations, "Overlap between expanded and optimized region not supported"
             first_set, second_set = regions, expanded_regions