Fallback to ort if hadamard shapes not available

Xilinx · Oct 29, 2024 · 53bac08 · 53bac08
1 parent 6ad16bf
commit 53bac08
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/src/brevitas/graph/equalize.py b/src/brevitas/graph/equalize.py
@@ -46,7 +46,13 @@
 else:
     RMSNorm = object
 
-__all__ = ['GraphActivationEqualization', 'LayerwiseActivationEqualization', 'EqualizeGraph']
+__all__ = [
+    'GraphActivationEqualization',
+    'LayerwiseActivationEqualization',
+    'EqualizeGraph',
+    'LayerwiseActivationRotation',
+    'MergeLnAffine',
+    'GraphRotationEqualization']
 
 EPSILON = 1e-9
 
@@ -1285,6 +1291,12 @@ def _apply_rotate(model: nn.Module, regions: List[Region], full_rotation_method=
                 rot_mat, K = get_hadK(hidden_dim)
             except AssertionError as e:
                 print(f"Incomptible shapes {hidden_dim}")
+                if not insert_rotation_module:
+                    print("Falling back to orthogonal matrices")
+                    rot_mat = random_orthogonal_matrix(hidden_dim)
+                    K = None
+                    rot_func = _apply_ort_device
+                print("Skipping layers")
                 continue
             rot_func = _apply_had_device