LSTM from Pytorch to Flax NNX

[JavierJJR]

Hi 👋,
I'm new to using Jax and I'm finding it challenging to convert my models from PyTorch to Flax NNX, particularly the LSTMs. For example, the following model processes data in batches of length 64 with 553 timesteps and 1 feature (64, 553, 1). How would the same model be implemented in Flax NNX? I feel that it would be helpful to supplement the package's documentation with more examples to assist new users.

class RNN_Model(nn.Module):
    def __init__(self, input_size=1, hidden_size=128, num_layers=2, output_size=10):
        super(RNN_Model, self).__init__()

        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size,
                            num_layers=num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        lstm_out, _ = self.lstm(x) 
        out = self.fc(lstm_out[:, -1, :])  
        return out