Implement RNN, GRU, LSTM (#268)

* RNN base implementation * Address comments+format * nits in docs * add tests for prb * fix test * add a couple tests --------- Co-authored-by: Awni Hannun <[email protected]>
2024-03-11 21:14:44 -07:00
parent 0e95b64942
commit 8e5600022a
6 changed files with 361 additions and 1 deletions
@@ -21,9 +21,11 @@ Layers
   Embedding
   GELU
   GroupNorm
+   GRU
   InstanceNorm
   LayerNorm
   Linear
+   LSTM
   MaxPool1d
   MaxPool2d
   Mish
@@ -32,6 +34,7 @@ Layers
   QuantizedLinear
   RMSNorm
   ReLU
+   RNN
   RoPE
   SELU
   Sequential