added MAMBA architecture

nanodl/__init__.py

@@ -107,6 +107,12 @@
     UNetResidualBlock
 )
 
+from nanodl.__src.models.mamba import (
+    Mamba,
+    MambaDataParallelTrainer,
+    MambaBlock
+)
+
 from nanodl.__src.models.transformer import (
     Transformer,
     TransformerDataParallelTrainer,
@@ -231,6 +237,9 @@
     "Mixtral",
     "MixtralDecoder",
     "MixtralDecoderBlock",
+    "Mamba",
+    "MambaDataParallelTrainer",
+    "MambaBlock"
     "Whisper",
     "WhisperDataParallelTrainer",
     "WhisperSpeechEncoder",
@@ -321,11 +330,15 @@ def test_jax(jax):
 def test_optax(optax):
     optimizer = optax.sgd(learning_rate=0.1)
 
+def test_einops(einops):
+    arr = einops.rearrange([1, 2, 3], 'a b c -> b a c')
+
 def main():
     try:
         flax = check_library_installed('flax')
         jax = check_library_installed('jax')
         optax = check_library_installed('optax')
+        einops = check_library_installed('einops')
 
         test_flax(flax)
         test_jax(jax)

nanodl/__src/models/gemma.py

@@ -307,9 +307,7 @@ def __call__(self,
 
 class Gemma(nn.Module):
     """
-    Implements the LLaMA2 model for text generation, featuring grouped rotary positional embeddings.
-
-    LLaMA2 enhances the transformer architecture by incorporating grouped rotary positional embeddings within its decoder blocks, aiming to improve the model's understanding of positional context and its ability to generate coherent and contextually relevant text.
+    Implements the Gemma model for text generation, featuring GQA + RMSNorm + RoPE.
 
     Attributes:
         num_layers (int): Number of layers (blocks) in the LLaMA2 model.