reformat and change Tensor_Parallel_Linear to Linear

axonn-ai · Oct 17, 2023 · 34abe16 · 34abe16
1 parent 33caaf3
commit 34abe16
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/axonn/intra_layer/__init__.py b/axonn/intra_layer/__init__.py
@@ -1,4 +1,4 @@
-from .fully_connected import Linear as Tensor_Parallel_Linear  # noqa: F401
+from .fully_connected import Linear  # noqa: F401
 from .communication import Drop, Gather
 from axonn import axonn as ax
 

diff --git a/axonn/intra_layer/fully_connected.py b/axonn/intra_layer/fully_connected.py
@@ -108,7 +108,9 @@ def forward(self, x, scatter_input=True, gather_output=True):
 
         bias = self.bias
         if gather_output:
-            bias = Gather.apply(self.bias, self.outer_group if not self.transpose else self.inner_group)
+            bias = Gather.apply(
+                self.bias, self.outer_group if not self.transpose else self.inner_group
+            )
 
         if self.skip_bias_add:
             return x, bias