renamed

tobiaskatsch · Feb 21, 2024 · af98340 · af98340
1 parent b00a32e
commit af98340
Show file tree

Hide file tree

Showing 17 changed files with 17 additions and 16 deletions.
diff --git a/README.md b/README.md
@@ -32,7 +32,7 @@ Other requirements:
 
 ## Usage
 We provide 2 main modules:
-- ### [gated_linear_rnn.py](gated_linear_rnn/gated_linear_rnn.py)
+- ### [gated_linear_rnn.py](flax_gated_linear_rnn/gated_linear_rnn.py)
   A causal time mixing sequence model which can be used as a drop-in replacement for causal multi-head-attention.
   Usage:
   ```
@@ -69,7 +69,7 @@ We provide 2 main modules:
   - **Tied Input & Forget gate** (`use_tied_gates=True`) Ties the input and forget gate through the relation `forget_gate = 1-input_gate`.
 
 
-- ## [gated_linear_rnn_lm.py](gated_linear_rnn/language_models/gated_linear_rnn_lm.py)
+- ## [gated_linear_rnn_lm.py](flax_gated_linear_rnn/language_models/gated_linear_rnn_lm.py)
   A GatedLinearRNN-based language model.
   ```
   import jax

diff --git a/flax_gated_linear_rnn/__init__.py b/flax_gated_linear_rnn/__init__.py
@@ -0,0 +1,3 @@
+from .gated_linear_rnn import GatedLinearRNN
+from .language_models import GatedLinearRNNLM
+
diff --git a/gated_linear_rnn/attention.py → flax_gated_linear_rnn/attention.py b/gated_linear_rnn/attention.py → flax_gated_linear_rnn/attention.py
diff --git a/gated_linear_rnn/base_models/__init__.py → ..._gated_linear_rnn/base_models/__init__.py b/gated_linear_rnn/base_models/__init__.py → ..._gated_linear_rnn/base_models/__init__.py
diff --git a/..._linear_rnn/base_models/channel_mixing.py → ..._linear_rnn/base_models/channel_mixing.py b/..._linear_rnn/base_models/channel_mixing.py → ..._linear_rnn/base_models/channel_mixing.py
diff --git a/..._linear_rnn/base_models/sequence_model.py → ..._linear_rnn/base_models/sequence_model.py b/..._linear_rnn/base_models/sequence_model.py → ..._linear_rnn/base_models/sequence_model.py
@@ -1,6 +1,6 @@
 from flax import linen as nn
 import jax.numpy as jnp
-from gated_linear_rnn.base_models.channel_mixing import ChannelMixing
+from flax_gated_linear_rnn.base_models.channel_mixing import ChannelMixing
 
 class SequenceModel(nn.Module):
     n_layer: int

diff --git a/gated_linear_rnn/base_models/time_mixing.py → ...ted_linear_rnn/base_models/time_mixing.py b/gated_linear_rnn/base_models/time_mixing.py → ...ted_linear_rnn/base_models/time_mixing.py
diff --git a/gated_linear_rnn/gated_linear_rnn.py → flax_gated_linear_rnn/gated_linear_rnn.py b/gated_linear_rnn/gated_linear_rnn.py → flax_gated_linear_rnn/gated_linear_rnn.py
diff --git a/gated_linear_rnn/language_models/__init__.py → ...ed_linear_rnn/language_models/__init__.py b/gated_linear_rnn/language_models/__init__.py → ...ed_linear_rnn/language_models/__init__.py
diff --git a/...nn/language_models/gated_linear_rnn_lm.py → ...nn/language_models/gated_linear_rnn_lm.py b/...nn/language_models/gated_linear_rnn_lm.py → ...nn/language_models/gated_linear_rnn_lm.py
@@ -1,7 +1,7 @@
-from gated_linear_rnn.base_models.sequence_model import SequenceModel
-from gated_linear_rnn.base_models.time_mixing import CausalTimeMixing
+from flax_gated_linear_rnn.base_models.sequence_model import SequenceModel
+from flax_gated_linear_rnn.base_models.time_mixing import CausalTimeMixing
 from typing import Optional, Callable
-from gated_linear_rnn.gated_linear_rnn import GatedLinearRNN
+from flax_gated_linear_rnn.gated_linear_rnn import GatedLinearRNN
 from flax import linen as nn
 
 class GatedLinearRNNLM(SequenceModel):

diff --git a/...ear_rnn/language_models/transformer_lm.py → ...ear_rnn/language_models/transformer_lm.py b/...ear_rnn/language_models/transformer_lm.py → ...ear_rnn/language_models/transformer_lm.py
@@ -1,7 +1,7 @@
 from typing import Optional, Callable
-from gated_linear_rnn.base_models.sequence_model import SequenceModel
-from gated_linear_rnn.base_models.time_mixing import CausalTimeMixing
-from gated_linear_rnn.attention import MultiHeadSelfAttention
+from flax_gated_linear_rnn.base_models.sequence_model import SequenceModel
+from flax_gated_linear_rnn.base_models.time_mixing import CausalTimeMixing
+from flax_gated_linear_rnn.attention import MultiHeadSelfAttention
 
 
 class TransformerLM(SequenceModel):

diff --git a/gated_linear_rnn/text_2_speech_models.py → ..._gated_linear_rnn/text_2_speech_models.py b/gated_linear_rnn/text_2_speech_models.py → ..._gated_linear_rnn/text_2_speech_models.py
diff --git a/gated_linear_rnn/__init__.py b/gated_linear_rnn/__init__.py
diff --git a/run_experiment.py b/run_experiment.py
@@ -4,8 +4,8 @@
 import ast
 from datetime import datetime
 import wandb
-from gated_linear_rnn import GateLoopLM, GateLoopText2SpeechModel
-from gated_linear_rnn.language_models import TransformerLM
+from flax_gated_linear_rnn import GateLoopLM, GateLoopText2SpeechModel
+from flax_gated_linear_rnn.language_models import TransformerLM
 
 
 

diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='flax-gated-linear-rnn',
-    version='1.0.0',
+    version='1.0.1',
     author='Tobias Katsch',
     author_email='tobias.katsch42@gmail.com',
     packages=find_packages(),

diff --git a/utils/speech_util.py b/utils/speech_util.py
@@ -8,7 +8,7 @@
 import torchaudio
 import numpy as np
 import os
-from gated_linear_rnn import GateLoopLM, GateLoopText2SpeechModel
+from flax_gated_linear_rnn import GateLoopLM, GateLoopText2SpeechModel
 
 
 vocab = {

diff --git a/utils/util.py b/utils/util.py
@@ -1,7 +1,7 @@
 import os
 from flax.training import train_state
 from typing import Any, Sequence, Optional, Tuple, Iterator, Dict, Callable, List
-from gated_linear_rnn.gated_linear_rnn import *
+from flax_gated_linear_rnn.gated_linear_rnn import *
 
 
 def get_home_directory():