[CLEANUP]

bit_moe_example.py

@@ -11,4 +11,4 @@
 output = model(x)
 
 # Print the output
-print(output)
+print(output)

bitnet/bit_moe.py

@@ -3,6 +3,9 @@
 from bitnet.bitlinear import BitLinear
 import torch.nn.functional as F
 
+# Feedforward = mlp = expert
+# = Linear projection + non linear activation functions like [RELU, GELU, etc] + Dropout[optional] + Normalization[optional, LayerNorm]
+
 
 # Expert module
 class Expert(nn.Module):
@@ -86,6 +89,7 @@ class BitMoE(nn.Module):
         num_experts (int): The number of experts in the mixture.
         top_k (int, optional): The number of experts to select for each input. Defaults to 2.
     """
+
     def __init__(self, dim: int, num_experts: int, top_k: int = 2):
         super(BitMoE, self).__init__()
         self.router = NoisyTopkRouter(dim, num_experts, top_k)
@@ -123,4 +127,4 @@ def forward(self, x):
 # x = torch.randn(2, 4, 8)
 # model = BitMoE(8, 4, 2)
 # output = model(x)
-# print(output)
+# print(output)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "bitnet"
-version = "0.2.4"
+version = "0.2.5"
 description = "bitnet - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
@@ -22,7 +22,7 @@ classifiers = [
 ]
 
 [tool.poetry.dependencies]
-python = "^3.6"
+python = "^3.10"
 torch = "*"
 einops = "*"
 zetascale = "*"