Merge pull request #4272 from zinccat:rnn_nnx

PiperOrigin-RevId: 686670119
google · Oct 16, 2024 · 3bf732c · 3bf732c
2 parents ae5c662 + 8c86f8e
commit 3bf732c
Show file tree

Hide file tree

Showing 3 changed files with 1,473 additions and 0 deletions.
diff --git a/flax/nnx/__init__.py b/flax/nnx/__init__.py
@@ -86,6 +86,13 @@
 from .nn.attention import dot_product_attention as dot_product_attention
 from .nn.attention import make_attention_mask as make_attention_mask
 from .nn.attention import make_causal_mask as make_causal_mask
+from .nn.recurrent import RNNCellBase as RNNCellBase
+from .nn.recurrent import LSTMCell as LSTMCell
+from .nn.recurrent import GRUCell as GRUCell
+from .nn.recurrent import OptimizedLSTMCell as OptimizedLSTMCell
+from .nn.recurrent import SimpleCell as SimpleCell
+from .nn.recurrent import RNN as RNN
+from .nn.recurrent import Bidirectional as Bidirectional
 from .nn.linear import Conv as Conv
 from .nn.linear import ConvTranspose as ConvTranspose
 from .nn.linear import Embed as Embed