Why dont we use a single Dense layer in LSTM as e.g. dm-Haiku does

[leonl42]

In the flax LSTMCell implementation we use 8 different Dense layers. Using a single layer and splitting (as e.g. dm-Haiku does) is significantly faster. I got a ~1.7x speed-up. Is there any reason we don't do this here?

[cgarciae]

Have you check OptimizedLSTMCell?

[leonl42]

Ah thank you, I didnt see that. So I guess the normal LSTMCell is just for legacy reasons because the weights are stored different? Because to me it does not make sense to have both. Why would you have the unoptimized LSTM when theres an optimized version? Wouldn't it then make sense to add a DeprecationWarning and replace the normal LSTM with the optimized one in the future?

[cgarciae]

Its a good question. I'm not sure about the history, but I understand that LSTMCell is seen as a readable implementation while the optimized version is used for performance.